iter

elastic · Mar 25, 2024 · 4e3c3d2 · 4e3c3d2
1 parent d4bacde
commit 4e3c3d2
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 58 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/runtime/AbstractScriptFieldQuery.java b/server/src/main/java/org/elasticsearch/search/runtime/AbstractScriptFieldQuery.java
@@ -71,18 +71,7 @@ public boolean isCacheable(LeafReaderContext ctx) {
             public Scorer scorer(LeafReaderContext ctx) {
                 S scriptContext = scriptContextFunction.apply(ctx);
                 DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
-                TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
-                    @Override
-                    public boolean matches() {
-                        return AbstractScriptFieldQuery.this.matches(scriptContext, approximation.docID());
-                    }
-
-                    @Override
-                    public float matchCost() {
-                        return MATCH_COST;
-                    }
-                };
-                return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
+                return new ConstantScoreScorer(this, score(), scoreMode, createTwoPhaseIterator(scriptContext, approximation));
             }
 
             @Override
@@ -96,6 +85,20 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
         };
     }
 
+    protected TwoPhaseIterator createTwoPhaseIterator(S scriptContext, DocIdSetIterator approximation) {
+        return new TwoPhaseIterator(approximation) {
+            @Override
+            public boolean matches() {
+                return AbstractScriptFieldQuery.this.matches(scriptContext, approximation.docID());
+            }
+
+            @Override
+            public float matchCost() {
+                return MATCH_COST;
+            }
+        };
+    }
+
     protected abstract boolean matches(S scriptContext, int docId);
 
     @Override

diff --git a/...c/main/java/org/elasticsearch/search/runtime/AbstractStringScriptFieldAutomatonQuery.java b/...c/main/java/org/elasticsearch/search/runtime/AbstractStringScriptFieldAutomatonQuery.java
@@ -8,7 +8,9 @@
 
 package org.elasticsearch.search.runtime;
 
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.automaton.ByteRunAutomaton;
 import org.elasticsearch.script.Script;
@@ -17,7 +19,6 @@
 import java.util.List;
 
 public abstract class AbstractStringScriptFieldAutomatonQuery extends AbstractStringScriptFieldQuery {
-    private final ThreadLocal<BytesRefBuilder> bytesRefBuilderThreadLocal = ThreadLocal.withInitial(BytesRefBuilder::new);
     private final ByteRunAutomaton automaton;
 
     public AbstractStringScriptFieldAutomatonQuery(
@@ -31,9 +32,24 @@ public AbstractStringScriptFieldAutomatonQuery(
     }
 
     @Override
-    protected final boolean matches(List<String> values) {
+    protected TwoPhaseIterator createTwoPhaseIterator(StringFieldScript scriptContext, DocIdSetIterator approximation) {
+        BytesRefBuilder scratch = new BytesRefBuilder();
+        return new TwoPhaseIterator(approximation) {
+            @Override
+            public boolean matches() {
+                scriptContext.runForDoc(approximation.docID());
+                return AbstractStringScriptFieldAutomatonQuery.this.matches(scriptContext.getValues(), scratch);
+            }
+
+            @Override
+            public float matchCost() {
+                return MATCH_COST;
+            }
+        };
+    }
+
+    protected final boolean matches(List<String> values, BytesRefBuilder scratch) {
         for (String value : values) {
-            BytesRefBuilder scratch = bytesRefBuilderThreadLocal.get();
             scratch.copyChars(value);
             if (automaton.run(scratch.bytes(), 0, scratch.length())) {
                 return true;
@@ -42,6 +58,11 @@ protected final boolean matches(List<String> values) {
         return false;
     }
 
+    @Override
+    protected final boolean matches(List<String> values) {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public final void visit(QueryVisitor visitor) {
         if (visitor.acceptField(fieldName())) {

diff --git a/server/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldFuzzyQueryTests.java b/server/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldFuzzyQueryTests.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.search.runtime;
 
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.ThreadInterruptedException;
 import org.apache.lucene.util.automaton.ByteRunAutomaton;
 import org.elasticsearch.script.Script;
@@ -74,29 +75,30 @@ protected StringScriptFieldFuzzyQuery mutate(StringScriptFieldFuzzyQuery orig) {
     @Override
     public void testMatches() {
         StringScriptFieldFuzzyQuery query = StringScriptFieldFuzzyQuery.build(randomScript(), leafFactory, "test", "foo", 1, 0, false);
-        assertTrue(query.matches(List.of("foo")));
-        assertTrue(query.matches(List.of("foa")));
-        assertTrue(query.matches(List.of("foo", "bar")));
-        assertFalse(query.matches(List.of("bar")));
+        BytesRefBuilder scratch = new BytesRefBuilder();
+        assertTrue(query.matches(List.of("foo"), scratch));
+        assertTrue(query.matches(List.of("foa"), scratch));
+        assertTrue(query.matches(List.of("foo", "bar"), scratch));
+        assertFalse(query.matches(List.of("bar"), scratch));
         query = StringScriptFieldFuzzyQuery.build(randomScript(), leafFactory, "test", "foo", 0, 0, false);
-        assertTrue(query.matches(List.of("foo")));
-        assertFalse(query.matches(List.of("foa")));
+        assertTrue(query.matches(List.of("foo"), scratch));
+        assertFalse(query.matches(List.of("foa"), scratch));
         query = StringScriptFieldFuzzyQuery.build(randomScript(), leafFactory, "test", "foo", 2, 0, false);
-        assertTrue(query.matches(List.of("foo")));
-        assertTrue(query.matches(List.of("foa")));
-        assertTrue(query.matches(List.of("faa")));
-        assertFalse(query.matches(List.of("faaa")));
+        assertTrue(query.matches(List.of("foo"), scratch));
+        assertTrue(query.matches(List.of("foa"), scratch));
+        assertTrue(query.matches(List.of("faa"), scratch));
+        assertFalse(query.matches(List.of("faaa"), scratch));
     }
 
     public void testConcurrentMatches() {
         StringScriptFieldFuzzyQuery query = StringScriptFieldFuzzyQuery.build(randomScript(), leafFactory, "test", "foo", 1, 0, false);
         List<Future<?>> futures = new ArrayList<>();
         ExecutorService executorService = Executors.newFixedThreadPool(3);
         try {
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foo")))));
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foa")))));
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foo", "bar")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("bar")))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foo"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foa"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("foo", "bar"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("bar"), new BytesRefBuilder()))));
             for (Future<?> future : futures) {
                 try {
                     future.get();

diff --git a/server/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldRegexpQueryTests.java b/server/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldRegexpQueryTests.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.search.runtime;
 
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.ThreadInterruptedException;
 import org.apache.lucene.util.automaton.ByteRunAutomaton;
 import org.apache.lucene.util.automaton.Operations;
@@ -90,13 +91,14 @@ public void testMatches() {
             0,
             Operations.DEFAULT_DETERMINIZE_WORK_LIMIT
         );
-        assertTrue(query.matches(List.of("astuffb")));
-        assertFalse(query.matches(List.of("astuffB")));
-        assertFalse(query.matches(List.of("fffff")));
-        assertFalse(query.matches(List.of("ab")));
-        assertFalse(query.matches(List.of("aasdf")));
-        assertFalse(query.matches(List.of("dsfb")));
-        assertTrue(query.matches(List.of("astuffb", "fffff")));
+        BytesRefBuilder scratch = new BytesRefBuilder();
+        assertTrue(query.matches(List.of("astuffb"), scratch));
+        assertFalse(query.matches(List.of("astuffB"), scratch));
+        assertFalse(query.matches(List.of("fffff"), scratch));
+        assertFalse(query.matches(List.of("ab"), scratch));
+        assertFalse(query.matches(List.of("aasdf"), scratch));
+        assertFalse(query.matches(List.of("dsfb"), scratch));
+        assertTrue(query.matches(List.of("astuffb", "fffff"), scratch));
 
         StringScriptFieldRegexpQuery ciQuery = new StringScriptFieldRegexpQuery(
             randomScript(),
@@ -125,13 +127,13 @@ public void testConcurrentMatches() {
         List<Future<?>> futures = new ArrayList<>();
         ExecutorService executorService = Executors.newFixedThreadPool(3);
         try {
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("astuffB")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("fffff")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("ab")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("aasdf")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("dsfb")))));
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb", "fffff")))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("astuffB"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("fffff"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("ab"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("aasdf"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("dsfb"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb", "fffff"), new BytesRefBuilder()))));
             for (Future<?> future : futures) {
                 try {
                     future.get();

diff --git a/...r/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldWildcardQueryTests.java b/...r/src/test/java/org/elasticsearch/search/runtime/StringScriptFieldWildcardQueryTests.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.search.runtime;
 
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.ThreadInterruptedException;
 import org.apache.lucene.util.automaton.ByteRunAutomaton;
 import org.elasticsearch.script.Script;
@@ -58,18 +59,19 @@ protected StringScriptFieldWildcardQuery mutate(StringScriptFieldWildcardQuery o
     @Override
     public void testMatches() {
         StringScriptFieldWildcardQuery query = new StringScriptFieldWildcardQuery(randomScript(), leafFactory, "test", "a*b", false);
-        assertTrue(query.matches(List.of("astuffb")));
-        assertFalse(query.matches(List.of("Astuffb")));
-        assertFalse(query.matches(List.of("fffff")));
-        assertFalse(query.matches(List.of("a")));
-        assertFalse(query.matches(List.of("b")));
-        assertFalse(query.matches(List.of("aasdf")));
-        assertFalse(query.matches(List.of("dsfb")));
-        assertTrue(query.matches(List.of("astuffb", "fffff")));
+        BytesRefBuilder scratch = new BytesRefBuilder();
+        assertTrue(query.matches(List.of("astuffb"), scratch));
+        assertFalse(query.matches(List.of("Astuffb"), scratch));
+        assertFalse(query.matches(List.of("fffff"), scratch));
+        assertFalse(query.matches(List.of("a"), scratch));
+        assertFalse(query.matches(List.of("b"), scratch));
+        assertFalse(query.matches(List.of("aasdf"), scratch));
+        assertFalse(query.matches(List.of("dsfb"), scratch));
+        assertTrue(query.matches(List.of("astuffb", "fffff"), scratch));
 
         StringScriptFieldWildcardQuery ciQuery = new StringScriptFieldWildcardQuery(randomScript(), leafFactory, "test", "a*b", true);
-        assertTrue(ciQuery.matches(List.of("Astuffb")));
-        assertTrue(ciQuery.matches(List.of("astuffB", "fffff")));
+        assertTrue(ciQuery.matches(List.of("Astuffb"), scratch));
+        assertTrue(ciQuery.matches(List.of("astuffB", "fffff"), scratch));
 
     }
 
@@ -78,11 +80,11 @@ public void testConcurrentMatches() {
         List<Future<?>> futures = new ArrayList<>();
         ExecutorService executorService = Executors.newFixedThreadPool(3);
         try {
-            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("Astuffb")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("fffff")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("a")))));
-            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("b")))));
+            futures.add(executorService.submit(() -> assertTrue(query.matches(List.of("astuffb"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("Astuffb"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("fffff"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("a"), new BytesRefBuilder()))));
+            futures.add(executorService.submit(() -> assertFalse(query.matches(List.of("b"), new BytesRefBuilder()))));
             for (Future<?> future : futures) {
                 try {
                     future.get();