sketchy implementation

Signed-off-by: bowenlan-amzn <[email protected]>
opensearch-project · Jul 16, 2024 · 63f1cd4 · 63f1cd4
1 parent 839d7c9
commit 63f1cd4
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 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
@@ -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<Object> 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
@@ -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<List<Object>> 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<Object> terms = new ArrayList<>();
+ terms.addAll(getResponse.getField(termsLookup.path()).getValues());
  if (getResponse.isSourceEmpty() == false) { // extract terms only if the doc source exists
  List<Object> 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
@@ -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);
  }