Repository: foundationdb/fdb-record-layer
Issue: LuceneRecordCursor.getTerms missing PrefixQuery

Because of the way this code is structured, it needs to know every possible Lucene search node.


Changed files:
fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordCursor.java
fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursorTest.java
fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursor.java
fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneHighlighting.java
fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexTest.java

diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursor.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursor.java
index 59abf528..1e39d9e6 100644
--- a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursor.java
+++ b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursor.java
@@ -220,7 +220,7 @@ public class LuceneAutoCompleteResultCursor implements BaseCursor<IndexEntry> {
         if (timer != null) {
             timer.increment(LuceneEvents.Counts.LUCENE_SCAN_MATCHED_AUTO_COMPLETE_SUGGESTIONS, topDocs.scoreDocs.length);
         }
-        return createResults(searcher, topDocs, tokenSet, prefixToken);
+        return createResults(searcher, topDocs, tokenSet, prefixToken == null ? Collections.emptySet() : Collections.singleton(prefixToken));
     }
 
     private Set<String> getAllIndexedFieldNames(IndexReader indexReader) {
@@ -355,10 +355,10 @@ public class LuceneAutoCompleteResultCursor implements BaseCursor<IndexEntry> {
     protected RecordCursor<IndexEntry> createResults(IndexSearcher searcher,
                                                      TopDocs topDocs,
                                                      Set<String> queryTokens,
-                                                     @Nullable String prefixToken) {
+                                                     Set<String> prefixTokens) {
         return RecordCursor.flatMapPipelined(
                 outerContinuation -> scoreDocsFromLookup(searcher, topDocs),
-                (scoreDocAndRecord, innerContinuation) -> findIndexEntriesInRecord(scoreDocAndRecord, queryTokens, prefixToken, innerContinuation),
+                (scoreDocAndRecord, innerContinuation) -> findIndexEntriesInRecord(scoreDocAndRecord, queryTokens, prefixTokens, innerContinuation),
                 scoreDocAndRecord -> scoreDocAndRecord.rec.getPrimaryKey().pack(),
                 null,
                 1 // Use a pipeline size of 1 because the inner cursors don't do I/O and the outer cursor has its own pipelining
@@ -406,7 +406,7 @@ public class LuceneAutoCompleteResultCursor implements BaseCursor<IndexEntry> {
         }
     }
 
-    private RecordCursor<IndexEntry> findIndexEntriesInRecord(ScoreDocAndRecord scoreDocAndRecord, Set<String> queryTokens, @Nullable String prefixToken, @Nullable byte[] continuation) {
+    private RecordCursor<IndexEntry> findIndexEntriesInRecord(ScoreDocAndRecord scoreDocAndRecord, Set<String> queryTokens, Set<String> prefixTokens, @Nullable byte[] continuation) {
         // Extract the indexed fields from the document again
         final List<LuceneDocumentFromRecord.DocumentField> documentFields = LuceneDocumentFromRecord.getRecordFields(state.index.getRootExpression(), scoreDocAndRecord.rec)
                 .get(groupingKey == null ? TupleHelpers.EMPTY : groupingKey)
@@ -428,7 +428,7 @@ public class LuceneAutoCompleteResultCursor implements BaseCursor<IndexEntry> {
                 // matched terms
                 return null;
             }
-            String match = LuceneHighlighting.searchAllMaybeHighlight(documentField.getFieldName(), queryAnalyzer, text, queryTokens, prefixToken, true,
+            String match = LuceneHighlighting.searchAllMaybeHighlight(documentField.getFieldName(), queryAnalyzer, text, queryTokens, prefixTokens, true,
                     new LuceneScanQueryParameters.LuceneQueryHighlightParameters(highlight), null);
             if (match == null) {
                 // Text not found in this field
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneHighlighting.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneHighlighting.java
index 651d692b..12b055cb 100644
--- a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneHighlighting.java
+++ b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneHighlighting.java
@@ -45,6 +45,7 @@ import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -64,7 +65,7 @@ public class LuceneHighlighting {
     @SuppressWarnings("squid:S3776") // Cognitive complexity is too high. Candidate for later refactoring
     @Nullable
     static String searchAllMaybeHighlight(@Nonnull String fieldName, @Nonnull Analyzer queryAnalyzer, @Nonnull String text,
-                                          @Nonnull Set<String> matchedTokens, @Nullable String prefixToken,
+                                          @Nonnull Set<String> matchedTokens, @Nonnull Set<String> prefixTokens,
                                           boolean allMatchingRequired,
                                           @Nonnull LuceneScanQueryParameters.LuceneQueryHighlightParameters luceneQueryHighlightParameters,
                                           @Nullable List<Pair<Integer, Integer>> highlightedPositions) {
@@ -75,7 +76,7 @@ public class LuceneHighlighting {
             StringBuilder sb = luceneQueryHighlightParameters.isHighlight() ? new StringBuilder() : null;
             int upto = 0;
             Set<String> matchedInText = new HashSet<>();
-            boolean matchedPrefix = false;
+            Set<String> matchedPrefixes = new HashSet<>();
             ArrayDeque<String> pres = new ArrayDeque<>();
             ArrayDeque<String> ends = new ArrayDeque<>();
             int lastMatchPos = -tokenCountAfterHighlighted - 1;
@@ -142,25 +143,30 @@ public class LuceneHighlighting {
                     upto = endOffset;
                     matchedInText.add(token);
                     lastMatchPos = currentPos;
-                } else if (prefixToken != null && token.startsWith(prefixToken)) {
-                    if (luceneQueryHighlightParameters.isHighlight()) {
-                        if (!tokenAlreadyHighlighted(text, startOffset, endOffset,
-                                luceneQueryHighlightParameters.getLeftTag(), luceneQueryHighlightParameters.getRightTag())) {
-                            addPrefixMatch(sb, text.substring(startOffset, endOffset), prefixToken,
-                                    luceneQueryHighlightParameters.getLeftTag(), luceneQueryHighlightParameters.getRightTag(),
-                                    highlightedPositions);
-                        } else {
-                            addNonMatch(sb, text.substring(startOffset, endOffset));
+                } else {
+                    for (String prefixToken : prefixTokens) {
+                        if (token.startsWith(prefixToken)) {
+                            if (luceneQueryHighlightParameters.isHighlight()) {
+                                if (!tokenAlreadyHighlighted(text, startOffset, endOffset,
+                                        luceneQueryHighlightParameters.getLeftTag(), luceneQueryHighlightParameters.getRightTag())) {
+                                    addPrefixMatch(sb, text.substring(startOffset, endOffset), prefixToken,
+                                            luceneQueryHighlightParameters.getLeftTag(), luceneQueryHighlightParameters.getRightTag(),
+                                            highlightedPositions);
+                                } else {
+                                    addNonMatch(sb, text.substring(startOffset, endOffset));
+                                }
+                            }
+                            upto = endOffset;
+                            matchedPrefixes.add(prefixToken);
+                            break;
                         }
                     }
-                    upto = endOffset;
-                    matchedPrefix = true;
                 }
                 currentPos++;
             }
             ts.end();
 
-            if (allMatchingRequired && ((prefixToken != null && !matchedPrefix) || (matchedInText.size() < matchedTokens.size()))) {
+            if (allMatchingRequired && (matchedPrefixes.size() < prefixTokens.size() || (matchedInText.size() < matchedTokens.size()))) {
                 // Query text not actually found in document text. Return null
                 return null;
             }
@@ -294,19 +300,20 @@ public class LuceneHighlighting {
                     if (terms.isEmpty()) {
                         return;
                     }
+                    Set<String> prefixes = getPrefixTerms(terms);
                     for (Map.Entry<Descriptors.FieldDescriptor, Object> entry : source.message.getAllFields().entrySet()) {
                         final Descriptors.FieldDescriptor entryDescriptor = entry.getKey();
                         final Object entryValue = entry.getValue();
                         if (entryValue instanceof String) {
                             buildIfMatch(source, fieldName, value,
                                     entryDescriptor, entryValue, 0,
-                                    terms, analyzerSelector, luceneQueryHighlightParameters);
+                                    terms, prefixes, analyzerSelector, luceneQueryHighlightParameters);
                         } else if (entryValue instanceof List) {
                             int index = 0;
                             for (Object entryValueElement : ((List<?>) entryValue)) {
                                 buildIfMatch(source, fieldName, value,
                                         entryDescriptor, entryValueElement, index,
-                                        terms, analyzerSelector, luceneQueryHighlightParameters);
+                                        terms, prefixes, analyzerSelector, luceneQueryHighlightParameters);
                                 index++;
                             }
                         }
@@ -314,17 +321,48 @@ public class LuceneHighlighting {
                 }, null);
     }
 
+    private static Set<String> getPrefixTerms(@Nonnull Set<String> terms) {
+        Set<String> result = Collections.emptySet();
+        Iterator<String> iter = terms.iterator();
+        while (iter.hasNext()) {
+            String term = iter.next();
+            if (term.endsWith("*")) {
+                term = term.substring(0, term.length() - 1);
+                if (result.isEmpty()) {
+                    result = new HashSet<>();
+                }
+                result.add(term);
+                iter.remove();
+            }
+        }
+        return result;
+    }
+
     private static <M extends Message> void buildIfMatch(RecordRebuildSource<M> source, String fieldName, Object fieldValue,
                                                          Descriptors.FieldDescriptor entryDescriptor, Object entryValue, int index,
-                                                         @Nonnull Set<String> terms,
+                                                         @Nonnull Set<String> terms, @Nonnull Set<String> prefixes,
                                                          @Nonnull LuceneAnalyzerCombinationProvider analyzerSelector,
                                                          @Nonnull LuceneScanQueryParameters.LuceneQueryHighlightParameters luceneQueryHighlightParameters) {
-        if (entryValue.equals(fieldValue) && terms.stream().anyMatch(t -> StringUtils.containsIgnoreCase((String)entryValue, t))) {
-            String highlightedText = searchAllMaybeHighlight(fieldName, analyzerSelector.provideIndexAnalyzer((String)entryValue).getAnalyzer(), (String)entryValue, terms, null, false, luceneQueryHighlightParameters, null);
+        if (entryValue.equals(fieldValue) && isMatch((String)entryValue, terms, prefixes)) {
+            String highlightedText = searchAllMaybeHighlight(fieldName, analyzerSelector.provideIndexAnalyzer((String)entryValue).getAnalyzer(), (String)entryValue, terms, prefixes, false, luceneQueryHighlightParameters, null);
             source.buildMessage(highlightedText, entryDescriptor, null, null, true, index);
         }
     }
 
+    private static boolean isMatch(@Nonnull String candidate, @Nonnull Set<String> terms, @Nonnull Set<String> prefixes) {
+        for (String term : terms) {
+            if (StringUtils.containsIgnoreCase(candidate, term)) {
+                return true;
+            }
+        }
+        for (String term : prefixes) {
+            if (StringUtils.containsIgnoreCase(candidate, term)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     static class RecordRebuildSource<M extends Message> implements LuceneIndexExpressions.RecordSource<RecordRebuildSource<M>> {
         @Nullable
         public final RecordRebuildSource<M> parent;
@@ -508,9 +546,10 @@ public class LuceneHighlighting {
                     if (terms.isEmpty()) {
                         return;
                     }
-                    if (value instanceof String && terms.stream().anyMatch(t -> StringUtils.containsIgnoreCase((String)value, t))) {
+                    Set<String> prefixes = getPrefixTerms(terms);
+                    if (value instanceof String && isMatch((String)value, terms, prefixes)) {
                         List<Pair<Integer, Integer>> highlightedPositions = new ArrayList<>();
-                        String highlightedText = searchAllMaybeHighlight(fieldName, analyzerSelector.provideIndexAnalyzer((String)value).getAnalyzer(), (String)value, terms, null, false, luceneQueryHighlightParameters, highlightedPositions);
+                        String highlightedText = searchAllMaybeHighlight(fieldName, analyzerSelector.provideIndexAnalyzer((String)value).getAnalyzer(), (String)value, terms, prefixes, false, luceneQueryHighlightParameters, highlightedPositions);
                         result.add(new HighlightedTerm(fieldName, highlightedText, highlightedPositions));
                     }
                 }, null);
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordCursor.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordCursor.java
index 29a9d53d..6c46839e 100644
--- a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordCursor.java
+++ b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordCursor.java
@@ -53,6 +53,7 @@ import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MultiPhraseQuery;
 import org.apache.lucene.search.PhraseQuery;
+import org.apache.lucene.search.PrefixQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.Sort;
@@ -434,6 +435,11 @@ class LuceneRecordCursor implements BaseCursor<IndexEntry> {
             Term term = spanTermQuery.getTerm();
             map.putIfAbsent(term.field(), new HashSet<>());
             map.get(term.field()).add(term.text().toLowerCase(Locale.ROOT));
+        } else if (query instanceof PrefixQuery) {
+            PrefixQuery termQuery = (PrefixQuery) query;
+            Term term = termQuery.getPrefix();
+            map.putIfAbsent(term.field(), new HashSet<>());
+            map.get(term.field()).add(term.text().toLowerCase(Locale.ROOT) + "*");
         } else {
             throw new RecordCoreException("This lucene query is not supported for highlighting");
         }
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursorTest.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursorTest.java
index 8b4415f4..b1d9ea0f 100644
--- a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursorTest.java
+++ b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneAutoCompleteResultCursorTest.java
@@ -126,7 +126,7 @@ class LuceneAutoCompleteResultCursorTest {
         assertEquals(expectedPrefixToken, prefixToken);
 
         Set<String> queryTokenSet = new HashSet<>(tokens);
-        @Nullable String match = LuceneHighlighting.searchAllMaybeHighlight("text", analyzer, text, queryTokenSet, prefixToken, true,
+        @Nullable String match = LuceneHighlighting.searchAllMaybeHighlight("text", analyzer, text, queryTokenSet, prefixToken == null ? Collections.emptySet() : Collections.singleton(prefixToken), true,
                 new LuceneScanQueryParameters.LuceneQueryHighlightParameters(highlight), null);
         assertEquals(expectedMatch, match);
     }
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexTest.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexTest.java
index fbe12aa4..d6222f33 100644
--- a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexTest.java
+++ b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexTest.java
@@ -778,6 +778,18 @@ public class LuceneIndexTest extends FDBRecordStoreTestBase {
         }
     }
 
+    @Test
+    void highlightedPrefix() {
+        try (FDBRecordContext context = openContext()) {
+            rebuildIndexMetaData(context, SIMPLE_DOC, SIMPLE_TEXT_SUFFIXES);
+            recordStore.saveRecord(createSimpleDocument(1645L, "Hello record layer", 1));
+            assertRecordTexts(List.of("Hello <b>recor</b>d layer"),
+                    recordStore.fetchIndexRecords(
+                            recordStore.scanIndex(SIMPLE_TEXT_SUFFIXES, fullTextSearch(SIMPLE_TEXT_SUFFIXES, "recor*", true), null, ScanProperties.FORWARD_SCAN),
+                            IndexOrphanBehavior.ERROR));
+        }
+    }
+
     @Test
     void testCommit() {
         try (FDBRecordContext context = openContext()) {