lazy load the score of point in set query

bowenlan-amzn · Oct 2, 2024 · a9d57b2 · a9d57b2
1 parent 8d19cb9
commit a9d57b2
Showing 1 changed file with 82 additions and 22 deletions.
diff --git a/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java
@@ -42,15 +42,22 @@
 import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.apache.lucene.document.StoredField;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.sandbox.document.BigIntegerPoint;
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.ConstantScoreWeight;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.IndexSortSortedNumericDocValuesRangeQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.PointInSetQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
+import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.NumericUtils;
 import org.opensearch.common.Explicit;
@@ -1508,36 +1515,89 @@ public static Query unsignedLongRangeQuery(
             return builder.apply(l, u);
         }
 
-        static PointInSetQuery bitmapIndexQuery(String field, RoaringBitmap bitmap) {
-            final BytesRef encoded = new BytesRef(new byte[Integer.BYTES]);
-            return new PointInSetQuery(field, 1, Integer.BYTES, new PointInSetQuery.Stream() {
+        static Query bitmapIndexQuery(String field, RoaringBitmap bitmap) {
+            return new Query() {
 
-                final Iterator<Integer> iterator = bitmap.iterator();
+                @Override
+                public String toString(String field) {
+                    return "";
+                }
 
                 @Override
-                public BytesRef next() {
-                    int value;
-                    if (iterator.hasNext()) {
-                        value = iterator.next();
-                    } else {
-                        return null;
-                    }
-                    IntPoint.encodeDimension(value, encoded.bytes, 0);
-                    return encoded;
+                public void visit(QueryVisitor visitor) {
+
                 }
-            }) {
+
                 @Override
-                public Query rewrite(IndexSearcher indexSearcher) throws IOException {
-                    if (bitmap.isEmpty()) {
-                        return new MatchNoDocsQuery();
-                    }
-                    return super.rewrite(indexSearcher);
+                public boolean equals(Object obj) {
+                    return false;
                 }
 
                 @Override
-                protected String toString(byte[] value) {
-                    assert value.length == Integer.BYTES;
-                    return Integer.toString(IntPoint.decodeDimension(value, 0));
+                public int hashCode() {
+                    return 0;
+                }
+
+                @Override
+                public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
+                    return new ConstantScoreWeight(this, boost) {
+                        @Override
+                        public Scorer scorer(LeafReaderContext context) throws IOException {
+                            return scorerSupplier(context).get(Long.MAX_VALUE);
+                        }
+
+                        @Override
+                        public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
+                            return new ScorerSupplier() {
+                                @Override
+                                public Scorer get(long leadCost) throws IOException {
+
+                                    final BytesRef encoded = new BytesRef(new byte[Integer.BYTES]);
+                                    Query query = new PointInSetQuery(field, 1, Integer.BYTES, new PointInSetQuery.Stream() {
+
+                                        final Iterator<Integer> iterator = bitmap.iterator();
+
+                                        @Override
+                                        public BytesRef next() {
+                                            int value;
+                                            if (iterator.hasNext()) {
+                                                value = iterator.next();
+                                            } else {
+                                                return null;
+                                            }
+                                            IntPoint.encodeDimension(value, encoded.bytes, 0);
+                                            return encoded;
+                                        }
+                                    }) {
+                                        @Override
+                                        public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+                                            if (bitmap.isEmpty()) {
+                                                return new MatchNoDocsQuery();
+                                            }
+                                            return super.rewrite(indexSearcher);
+                                        }
+
+                                        @Override
+                                        protected String toString(byte[] value) {
+                                            assert value.length == Integer.BYTES;
+                                            return Integer.toString(IntPoint.decodeDimension(value, 0));
+                                        }
+                                    };
+                                    return query.createWeight(searcher, scoreMode, boost).scorer(context);
+                                }
+
+                                @Override
+                                public long cost() {
+                                    return bitmap.getLongCardinality();
+                                }
+                            };
+                        }
+
+                        @Override
+                        public boolean isCacheable(LeafReaderContext ctx) {
+                            return false;
+                        }
+                    };
                 }
             };
         }