Sort terms once when terms query is created for indexed and docValues keywords field.

Signed-off-by: Mikhail Khludnev <[email protected]>

Author: mkhludnev

server/src/main/java/org/opensearch/index/mapper/BytesRefsCollectionBuilder.java

@@ -24,8 +24,8 @@
 
 /**
  * Purposed for passing terms into {@link TermInSetQuery}.
- * If the given terms are sorted already, it wrap it with a SortedSet stub.
- * Otherwise, it passes terms as list.
+ * If the given terms are sorted already, it wraps terms with a SortedSet stub.
+ * Otherwise, it passes terms as a list.
  */
 public class BytesRefsCollectionBuilder implements Consumer<BytesRef>, Supplier<Collection<BytesRef>> {
 

server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java

@@ -500,22 +500,41 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
                     return super.termsQuery(values, context);
                 }
                 BytesRefsCollectionBuilder iBytesRefs = new BytesRefsCollectionBuilder(values.size());
-                BytesRefsCollectionBuilder dVByteRefs = new BytesRefsCollectionBuilder(values.size());
+                BytesRefsCollectionBuilder dVByteRefs = null; // = new BytesRefsCollectionBuilder(values.size());
                 for (int i = 0; i < values.size(); i++) {
-                    BytesRef idxBytes = indexedValueForSearch(values.get(i));
+                    Object value = values.get(i);
+                    BytesRef idxBytes = indexedValueForSearch(value);
                     iBytesRefs.accept(idxBytes);
-                    BytesRef dvBytes = indexedValueForSearch(rewriteForDocValue(values.get(i)));
-                    dVByteRefs.accept(dvBytes);
+
+                    Object rewritten = rewriteForDocValue(value);
+                    if (dVByteRefs == null) { // need check
+                        if (rewritten != value && !rewritten.equals(value)) {
+                            // first time see the difference between index and doc values, let's clone
+                            dVByteRefs = new BytesRefsCollectionBuilder(values.size());
+                            for (int rewind = 0; rewind <= i; rewind++) {
+                                Object rewrittenOld = rewind < i ? rewriteForDocValue(values.get(rewind)) : rewritten;
+                                BytesRef dvBytesOld = indexedValueForSearch(rewrittenOld);
+                                dVByteRefs.accept(dvBytesOld);
+                            }
+                        }
+                    } else {
+                        BytesRef dvBytes = indexedValueForSearch(rewritten);
+                        dVByteRefs.accept(dvBytes);
+                    }
+                }
+                if (dVByteRefs == null) { // index and docValues are the same, pack them once
+                    return TermInSetQuery.newIndexOrDocValuesQuery(MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, name(), iBytesRefs.get());
+                } else {
+                    Query indexQuery = new TermInSetQuery(MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, name(), iBytesRefs.get());
+                    Query dvQuery = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), dVByteRefs.get());
+                    return new IndexOrDocValuesQuery(indexQuery, dvQuery);
                 }
-                Query indexQuery = new TermInSetQuery(MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, name(), iBytesRefs.get());
-                Query dvQuery = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), dVByteRefs.get());
-                return new IndexOrDocValuesQuery(indexQuery, dvQuery);
             }
             // if we only have doc_values enabled, we construct a new query with doc_values re-written
             if (hasDocValues()) {
                 BytesRefsCollectionBuilder bytesCollector = new BytesRefsCollectionBuilder(values.size());
-                for (int i = 0; i < values.size(); i++) {
-                    BytesRef dvBytes = indexedValueForSearch(rewriteForDocValue(values.get(i)));
+                for (Object value : values) {
+                    BytesRef dvBytes = indexedValueForSearch(rewriteForDocValue(value));
                     bytesCollector.accept(dvBytes);
                 }
                 return new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesCollector.get());

server/src/test/java/org/opensearch/index/mapper/KeywordFieldTypeTests.java

@@ -229,6 +229,35 @@ public void testTermsSortedQuery() {
         assertEquals(expectedDocValues, onlyDocValues.termsQuery(sortedStrings, null));
     }
 
+    public void testTermsQuery_WhenDocValuesTrickyRewrite() {
+        MappedFieldType ft = new KeywordFieldType("field") {
+            @Override
+            protected Object rewriteForDocValue(Object value) {
+                String strVal = (String) value;
+                if (strVal.compareTo("a") >= 0) { // prepend letters in DV with underscore
+                    return "_" + strVal;
+                } else {
+                    if (randomBoolean()) {
+                        return super.rewriteForDocValue(value);
+                    } else {
+                        return value;
+                    }
+                }
+            }
+        };
+        List<String> strings = List.of("1", "2", "3", "4", "5", "a", "b", "c", "d", "e");
+        List<String> dvStrings = List.of("1", "2", "3", "4", "5", "_a", "_b", "_c", "_d", "_e");
+        Query expected = new IndexOrDocValuesQuery(
+            new TermInSetQuery("field", strings.stream().map(BytesRef::new).toList()),
+            new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, "field", dvStrings.stream().map(BytesRef::new).toList())
+        );
+        if (rarely()) {
+            strings = new ArrayList<>(strings);
+            Collections.shuffle(strings, random());
+        }
+        assertEquals(expected, ft.termsQuery(strings, MOCK_QSC_ENABLE_INDEX_DOC_VALUES));
+    }
+
     public void testExistsQuery() {
         {
             KeywordFieldType ft = new KeywordFieldType("field");