From 0d672106506f07386efbc9430026e1f394248907 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christoph=20B=C3=BCscher?= <christophbuescher@posteo.de>
Date: Wed, 28 Aug 2024 14:39:43 +0200
Subject: [PATCH] Implement missing prepareSeekExact methods for TermsEnum

In https://github.com/apache/lucene/pull/13359 a new "prepareSeekExact" method
was added that can improve seeking on TermsEnum implementations. Two of our own
subclasses of TermsEnum don't seem to support seeking for text, so we can safely
throw an UOE there. The third (FilterableTermsEnum) changes to simple returning
a Supplier for the actual "seek" method for now.
---
 .../common/lucene/index/FilterableTermsEnum.java          | 8 +++++++-
 .../index/mapper/flattened/FlattenedFieldMapper.java      | 6 ++++++
 .../xpack/core/termsenum/action/SimpleTermCountEnum.java  | 6 ++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/elasticsearch/common/lucene/index/FilterableTermsEnum.java b/server/src/main/java/org/elasticsearch/common/lucene/index/FilterableTermsEnum.java
index e9ecfbad42811..3bc1949cdfcbf 100644
--- a/server/src/main/java/org/elasticsearch/common/lucene/index/FilterableTermsEnum.java
+++ b/server/src/main/java/org/elasticsearch/common/lucene/index/FilterableTermsEnum.java
@@ -26,6 +26,7 @@
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.IOBooleanSupplier;
 import org.elasticsearch.core.Nullable;
 
 import java.io.IOException;
@@ -61,7 +62,7 @@ static class Holder {
     protected BytesRef current;
     protected final int docsEnumFlag;
 
-    public FilterableTermsEnum(IndexReader reader, String field, int docsEnumFlag, @Nullable Query filter) throws IOException {
+    public FilterableTermsEnum(IndexReader reader, String field, int docsEnumFlag, @Nullable Query f) throws IOException {
         if ((docsEnumFlag != PostingsEnum.FREQS) && (docsEnumFlag != PostingsEnum.NONE)) {
             throw new IllegalArgumentException("invalid docsEnumFlag of " + docsEnumFlag);
         }
@@ -176,6 +177,11 @@ public boolean seekExact(BytesRef text) throws IOException {
         }
     }
 
+    @Override
+    public IOBooleanSupplier prepareSeekExact(BytesRef bytesRef) {
+        return () -> this.seekExact(bytesRef);
+    }
+
     @Override
     public int docFreq() throws IOException {
         return currentDocFreq;
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/flattened/FlattenedFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/flattened/FlattenedFieldMapper.java
index 0c1340867849a..986193880a123 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/flattened/FlattenedFieldMapper.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/flattened/FlattenedFieldMapper.java
@@ -27,6 +27,7 @@
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.util.AttributeSource;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.IOBooleanSupplier;
 import org.apache.lucene.util.automaton.Automata;
 import org.apache.lucene.util.automaton.Automaton;
 import org.apache.lucene.util.automaton.CompiledAutomaton;
@@ -459,6 +460,11 @@ public AttributeSource attributes() {
             throw new UnsupportedOperationException();
         }
 
+        @Override
+        public IOBooleanSupplier prepareSeekExact(BytesRef bytesRef) throws IOException {
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         public boolean seekExact(BytesRef text) throws IOException {
             throw new UnsupportedOperationException();
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/SimpleTermCountEnum.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/SimpleTermCountEnum.java
index b4952373dfdd3..92568c4f31c18 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/SimpleTermCountEnum.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/SimpleTermCountEnum.java
@@ -13,6 +13,7 @@
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.util.AttributeSource;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.IOBooleanSupplier;
 import org.elasticsearch.index.mapper.MappedFieldType;
 
 import java.io.IOException;
@@ -69,6 +70,11 @@ public AttributeSource attributes() {
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public IOBooleanSupplier prepareSeekExact(BytesRef bytesRef) throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public boolean seekExact(BytesRef text) throws IOException {
         throw new UnsupportedOperationException();