First attempt at replacing SortedSetDocValues with BinaryDocValues in keyword mapper

x-pack/plugin/wildcard/src/main/java/org/elasticsearch/xpack/wildcard/mapper/BinaryDvConfirmedQuery.java → server/src/main/java/org/elasticsearch/index/mapper/BinaryDvConfirmedQuery.java

@@ -1,11 +1,13 @@
 /*
  * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.xpack.wildcard.mapper;
+package org.elasticsearch.index.mapper;
 
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DocValues;
@@ -35,7 +37,7 @@
  * match a provided approximation query which is key to getting good performance).
  */
 
-abstract class BinaryDvConfirmedQuery extends Query {
+public abstract class BinaryDvConfirmedQuery extends Query {
 
     protected final String field;
     protected final Query approxQuery;
@@ -53,6 +55,10 @@ public static Query fromAutomaton(Query approximation, String field, String matc
         return new BinaryDvConfirmedAutomatonQuery(approximation, field, matchPattern, automaton);
     }
 
+    public static Query fromAutomatonSingleValue(Query approximation, String field, String matchPattern, Automaton automaton) {
+        return new SingleValueBinaryDvConfirmedAutomatonQuery(approximation, field, matchPattern, automaton);
+    }
+
     /**
      * Returns a query that checks for equality of at leat one of the provided terms across
      * all binary doc values (but only for docs that also match a provided approximation query which
@@ -63,6 +69,11 @@ public static Query fromTerms(Query approximation, String field, BytesRef... ter
         return new BinaryDvConfirmedTermsQuery(approximation, field, terms);
     }
 
+    public static Query fromTermsSingleValue(Query approximation, String field, BytesRef... terms) {
+        Arrays.sort(terms, BytesRef::compareTo);
+        return new SingleValueBinaryDvConfirmedTermsQuery(approximation, field, terms);
+    }
+
     protected abstract boolean matchesBinaryDV(ByteArrayStreamInput bytes, BytesRef bytesRef, BytesRef scratch) throws IOException;
 
     protected abstract Query rewrite(Query approxRewrite) throws IOException;
@@ -146,7 +157,7 @@ public int hashCode() {
         return Objects.hash(classHash(), field, approxQuery);
     }
 
-    Query getApproximationQuery() {
+    public Query getApproximationQuery() {
         return approxQuery;
     }
 
@@ -159,7 +170,7 @@ public void visit(QueryVisitor visitor) {
 
     private static class BinaryDvConfirmedAutomatonQuery extends BinaryDvConfirmedQuery {
 
-        private final ByteRunAutomaton byteRunAutomaton;
+        protected final ByteRunAutomaton byteRunAutomaton;
         private final String matchPattern;
 
         private BinaryDvConfirmedAutomatonQuery(Query approximation, String field, String matchPattern, Automaton automaton) {
@@ -209,9 +220,20 @@ public int hashCode() {
         }
     }
 
+    private static class SingleValueBinaryDvConfirmedAutomatonQuery extends BinaryDvConfirmedAutomatonQuery {
+        private SingleValueBinaryDvConfirmedAutomatonQuery(Query approximation, String field, String matchPattern, Automaton automaton) {
+            super(approximation, field, matchPattern, automaton);
+        }
+
+        @Override
+        protected boolean matchesBinaryDV(ByteArrayStreamInput bytes, BytesRef bytesRef, BytesRef scratch) {
+            return byteRunAutomaton.run(bytesRef.bytes, bytesRef.offset, bytesRef.length);
+        }
+    }
+
     private static class BinaryDvConfirmedTermsQuery extends BinaryDvConfirmedQuery {
 
-        private final BytesRef[] terms;
+        protected final BytesRef[] terms;
 
         private BinaryDvConfirmedTermsQuery(Query approximation, String field, BytesRef[] terms) {
             super(approximation, field);
@@ -275,4 +297,24 @@ public int hashCode() {
             return Objects.hash(super.hashCode(), Arrays.hashCode(terms));
         }
     }
+
+    static class SingleValueBinaryDvConfirmedTermsQuery extends BinaryDvConfirmedTermsQuery {
+        SingleValueBinaryDvConfirmedTermsQuery(Query approximation, String field, BytesRef[] terms) {
+            super(approximation, field, terms);
+        }
+
+        @Override
+        protected boolean matchesBinaryDV(ByteArrayStreamInput bytes, BytesRef bytesRef, BytesRef scratch) {
+            if (terms.length == 1) {
+                return terms[0].bytesEquals(bytesRef);
+            } else {
+                final int pos = Arrays.binarySearch(terms, bytesRef, BytesRef::compareTo);
+                if (pos >= 0) {
+                    assert terms[pos].bytesEquals(bytesRef) : "Expected term at position " + pos + " to match bytesRef, but it did not.";
+                    return true;
+                }
+                return false;
+            }
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

@@ -707,6 +707,33 @@ public String toString() {
         }
     }
 
+    public static class BytesRefsFromSingletonBinaryBlockLoader extends DocValuesBlockLoader {
+        private final String fieldName;
+
+        public BytesRefsFromSingletonBinaryBlockLoader(String fieldName) {
+            this.fieldName = fieldName;
+        }
+
+        @Override
+        public BytesRefBuilder builder(BlockFactory factory, int expectedCount) {
+            return factory.bytesRefs(expectedCount);
+        }
+
+        @Override
+        public AllReader reader(LeafReaderContext context) throws IOException {
+            BinaryDocValues docValues = context.reader().getBinaryDocValues(fieldName);
+            if (docValues != null) {
+                return new BytesRefsFromSingletonBinary(docValues);
+            }
+            return new ConstantNullsReader();
+        }
+
+        @Override
+        public String toString() {
+            return "BytesRefsFromSingletonBinaryBlockLoader[" + fieldName + "]";
+        }
+    }
+
     public static class BytesRefsFromOrdsBlockLoader extends DocValuesBlockLoader {
         private final String fieldName;
 
@@ -1005,8 +1032,8 @@ public String toString() {
      * Read BinaryDocValues with no additional structure in the BytesRefs.
      * Each BytesRef from the doc values maps directly to a value in the block loader.
      */
-    public static class BytesRefsFromBinary extends AbstractBytesRefsFromBinary {
-        public BytesRefsFromBinary(BinaryDocValues docValues) {
+    public static class BytesRefsFromSingletonBinary extends AbstractBytesRefsFromBinary {
+        public BytesRefsFromSingletonBinary(BinaryDocValues docValues) {
             super(docValues);
         }