1 files changed, 17 insertions, 11 deletions

diff --git a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
index f12e9b0676..4ad2e7824e 100644
--- a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
@@ -20,11 +20,14 @@
 package org.elasticsearch.index.query;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.ConstantScoreScorer;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.search.RandomAccessWeight;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.search.Weight;
-import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -172,16 +175,18 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
         }
 
         @Override
-        public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException {
-            return new RandomAccessWeight(this) {
+        public Weight createWeight(IndexSearcher searcher, boolean needsScores, float boost) throws IOException {
+            return new ConstantScoreWeight(this, boost) {
+
                 @Override
-                protected Bits getMatchingDocs(final LeafReaderContext context) throws IOException {
+                public Scorer scorer(LeafReaderContext context) throws IOException {
+                    DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
                     final LeafSearchScript leafScript = searchScript.getLeafSearchScript(context);
-                    return new Bits() {
+                    TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
 
                         @Override
-                        public boolean get(int doc) {
-                            leafScript.setDocument(doc);
+                        public boolean matches() throws IOException {
+                            leafScript.setDocument(approximation.docID());
                             Object val = leafScript.run();
                             if (val == null) {
                                 return false;
@@ -196,11 +201,12 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
                         }
 
                         @Override
-                        public int length() {
-                            return context.reader().maxDoc();
+                        public float matchCost() {
+                            // TODO: how can we compute this?
+                            return 1000f;
                         }
-
                     };
+                    return new ConstantScoreScorer(this, score(), twoPhase);
                 }
             };
         }