From 63f1cd4ec3eab6261b3ba854c5fd2dbc36cfd27e Mon Sep 17 00:00:00 2001 From: bowenlan-amzn Date: Tue, 16 Jul 2024 10:43:29 -0700 Subject: [PATCH] sketchy implementation Signed-off-by: bowenlan-amzn --- .../index/mapper/NumberFieldMapper.java | 16 +++-- .../index/query/TermsQueryBuilder.java | 11 +--- .../search/query/BitMapFilterQuery.java | 64 ++++++++++--------- 3 files changed, 48 insertions(+), 43 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 index ea33bb17a91ee..d60defe2dbd9b 100644 --- a/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java @@ -58,6 +58,7 @@ import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Setting.Property; import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.core.xcontent.XContentParser.Token; import org.opensearch.index.document.SortedUnsignedLongDocValuesRangeQuery; @@ -69,10 +70,10 @@ import org.opensearch.search.DocValueFormat; import org.opensearch.search.lookup.SearchLookup; import org.opensearch.search.query.BitMapFilterQuery; -import org.roaringbitmap.RoaringBitmap; import java.io.IOException; import java.math.BigInteger; +import java.nio.ByteBuffer; import java.time.ZoneId; import java.util.ArrayList; import java.util.Arrays; @@ -84,6 +85,8 @@ import java.util.function.Function; import java.util.function.Supplier; +import org.roaringbitmap.RoaringBitmap; + /** * A {@link FieldMapper} for numeric types: byte, short, int, long, float and double. * @@ -790,9 +793,14 @@ public Query termQuery(String field, Object value, boolean hasDocValues, boolean @Override public Query termsQuery(String field, List values, boolean hasDocValues, boolean isSearchable) { - if (values.get(0) instanceof RoaringBitmap) { - RoaringBitmap bm = (RoaringBitmap) values.get(0); - return new BitMapFilterQuery(field, bm); + if (values.get(0) instanceof BytesArray) { + RoaringBitmap bitmap = new RoaringBitmap(); + try { + bitmap.deserialize(ByteBuffer.wrap(((BytesArray) values.get(0)).array())); + return new BitMapFilterQuery(field, bitmap); + } catch (IOException e) { + throw new RuntimeException(e); + } } int[] v = new int[values.size()]; diff --git a/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java index 75d3cda605fac..4f35fb20a91a8 100644 --- a/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java @@ -54,7 +54,6 @@ import org.opensearch.index.mapper.ConstantFieldType; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.indices.TermsLookup; -import org.roaringbitmap.RoaringBitmap; import java.io.IOException; import java.nio.CharBuffer; @@ -474,20 +473,16 @@ protected Query doToQuery(QueryShardContext context) throws IOException { if (fieldType == null) { throw new IllegalStateException("Rewrite first"); } - - RoaringBitmap bm = new RoaringBitmap(); - for (Object value : values) { - bm.add(((Long) value).intValue()); - } - - return fieldType.termsQuery(List.of(bm), context); + return fieldType.termsQuery(values, context); } private void fetch(TermsLookup termsLookup, Client client, ActionListener> actionListener) { GetRequest getRequest = new GetRequest(termsLookup.index(), termsLookup.id()); getRequest.preference("_local").routing(termsLookup.routing()); + getRequest.storedFields(termsLookup.path()); client.get(getRequest, ActionListener.delegateFailure(actionListener, (delegatedListener, getResponse) -> { List terms = new ArrayList<>(); + terms.addAll(getResponse.getField(termsLookup.path()).getValues()); if (getResponse.isSourceEmpty() == false) { // extract terms only if the doc source exists List extractedValues = XContentMapValues.extractRawValues(termsLookup.path(), getResponse.getSourceAsMap()); terms.addAll(extractedValues); diff --git a/server/src/main/java/org/opensearch/search/query/BitMapFilterQuery.java b/server/src/main/java/org/opensearch/search/query/BitMapFilterQuery.java index 417077be40c89..0dcdbe23f7188 100644 --- a/server/src/main/java/org/opensearch/search/query/BitMapFilterQuery.java +++ b/server/src/main/java/org/opensearch/search/query/BitMapFilterQuery.java @@ -22,10 +22,14 @@ import org.apache.lucene.search.TwoPhaseIterator; import org.apache.lucene.search.Weight; import org.apache.lucene.util.Accountable; -import org.roaringbitmap.RoaringBitmap; import java.io.IOException; +import org.roaringbitmap.RoaringBitmap; + +/** + * Filter with bitmap + */ public class BitMapFilterQuery extends Query implements Accountable { final String field; @@ -45,39 +49,37 @@ public Scorer scorer(LeafReaderContext context) throws IOException { final NumericDocValues singleton = DocValues.unwrapSingleton(values); final TwoPhaseIterator iterator; if (singleton != null) { - iterator = - new TwoPhaseIterator(singleton) { - @Override - public boolean matches() throws IOException { - long value = singleton.longValue(); - return bitmap.contains((int) value); - } - - @Override - public float matchCost() { - return 5; // 2 comparisions, possible lookup in the set - } - }; + iterator = new TwoPhaseIterator(singleton) { + @Override + public boolean matches() throws IOException { + long value = singleton.longValue(); + return bitmap.contains((int) value); + } + + @Override + public float matchCost() { + return 5; // 2 comparisions, possible lookup in the set + } + }; } else { - iterator = - new TwoPhaseIterator(values) { - @Override - public boolean matches() throws IOException { - int count = values.docValueCount(); - for (int i = 0; i < count; i++) { - final long value = values.nextValue(); - if (bitmap.contains((int) value)) { - return true; - } + iterator = new TwoPhaseIterator(values) { + @Override + public boolean matches() throws IOException { + int count = values.docValueCount(); + for (int i = 0; i < count; i++) { + final long value = values.nextValue(); + if (bitmap.contains((int) value)) { + return true; } - return false; - } - - @Override - public float matchCost() { - return 5; // 2 comparisons, possible lookup in the set } - }; + return false; + } + + @Override + public float matchCost() { + return 5; // 2 comparisons, possible lookup in the set + } + }; } return new ConstantScoreScorer(this, score(), scoreMode, iterator); }