diff --git a/CHANGELOG.md b/CHANGELOG.md index b37c3e4e74161..b8596d083c280 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 3.x] ### Added +- Add bitmap64 query support ([#20606](https://github.com/opensearch-project/OpenSearch/pull/20606)) - Add ProfilingWrapper interface for plugin access to delegates in profiling decorators ([#20607](https://github.com/opensearch-project/OpenSearch/pull/20607)) - Support expected cluster name with validation in CCS Sniff mode ([#20532](https://github.com/opensearch-project/OpenSearch/pull/20532)) - Choose the best performing node when writing with append-only index ([#20065](https://github.com/opensearch-project/OpenSearch/pull/20065)) diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search/381_bitmap_filtering_long.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search/381_bitmap_filtering_long.yml new file mode 100644 index 0000000000000..1607a3bc47922 --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/search/381_bitmap_filtering_long.yml @@ -0,0 +1,99 @@ +--- +setup: + - skip: + version: " - 3.5.99" + reason: Bitmap filtering for long fields is available in 3.6 and later. + + - do: + indices.create: + index: employees + body: + settings: + number_of_shards: 1 + number_of_replicas: 0 + mappings: + properties: + employee_id: + type: long + + - do: + bulk: + refresh: true + body: + - { "index": { "_index": "employees", "_id": "1" } } + - { "name": "Alice Smith", "employee_id": 1000000000001 } + - { "index": { "_index": "employees", "_id": "2" } } + - { "name": "Bob Johnson", "employee_id": 2000000000002 } + - { "index": { "_index": "employees", "_id": "3" } } + - { "name": "Charlie Brown", "employee_id": 3000000000003 } + + - do: + indices.create: + index: departments + body: + settings: + number_of_shards: 1 + number_of_replicas: 0 + mappings: + properties: + members: + type: binary + store: true + + - do: + bulk: + refresh: true + body: + - { "index": { "_index": "departments", "_id": "201" } } + - { "members": "AgAAAAAAAADoAAAAOjAAAAEAAACl1AAAEAAAAAEQ0QEAADowAAABAAAASqkAABAAAAACIA==" } + - { "index": { "_index": "departments", "_id": "202" } } + - { "members": "AQAAAAAAAADoAAAAOjAAAAEAAACl1AAAEAAAAAEQ" } + + - do: + cluster.health: + wait_for_status: green + +--- +"Terms lookup on a binary field with bitmap (long)": + - do: + search: + rest_total_hits_as_int: true + index: employees + body: { + "query": { + "terms": { + "employee_id": { + "index": "departments", + "id": "201", + "path": "members", + "store": true + }, + "value_type": "bitmap" + } + } + } + - match: { hits.total: 2 } + - match: { hits.hits.0._source.name: Alice Smith } + - match: { hits.hits.0._source.employee_id: 1000000000001 } + - match: { hits.hits.1._source.name: Bob Johnson } + - match: { hits.hits.1._source.employee_id: 2000000000002 } + +--- +"Terms query accepting bitmap as value (long)": + - do: + search: + rest_total_hits_as_int: true + index: employees + body: { + "query": { + "terms": { + "employee_id": ["AgAAAAAAAADoAAAAOjAAAAEAAACl1AAAEAAAAAEQ0QEAADowAAABAAAASqkAABAAAAACIA=="], + "value_type": "bitmap" + } + } + } + - match: { hits.total: 2 } + - match: { hits.hits.0._source.name: Alice Smith } + - match: { hits.hits.0._source.employee_id: 1000000000001 } + - match: { hits.hits.1._source.name: Bob Johnson } + - match: { hits.hits.1._source.employee_id: 2000000000002 } diff --git a/server/src/internalClusterTest/java/org/opensearch/search/query/SearchQueryIT.java b/server/src/internalClusterTest/java/org/opensearch/search/query/SearchQueryIT.java index cf4b0106f8090..55d3f41c3c104 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/query/SearchQueryIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/search/query/SearchQueryIT.java @@ -84,6 +84,8 @@ import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase; import org.opensearch.test.junit.annotations.TestIssueLogging; +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; import java.io.IOException; import java.io.Reader; import java.nio.ByteBuffer; @@ -102,6 +104,7 @@ import java.util.regex.Pattern; import org.roaringbitmap.RoaringBitmap; +import org.roaringbitmap.longlong.Roaring64NavigableMap; import static java.util.Collections.singletonMap; import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE; @@ -1197,6 +1200,88 @@ public void testTermsQueryWithBitmapDocValuesQuery() throws Exception { assertSearchHits(searchResponse, "1", "3", "4"); } + public void testTermsQueryWithBitmap64DocValuesQuery() throws Exception { + assertAcked( + prepareCreate("employees").setMapping( + jsonBuilder().startObject() + .startObject("properties") + .startObject("employee_id") + .field("type", "long") + .field("index", false) + .endObject() + .endObject() + .endObject() + ) + ); + indexRandom( + true, + client().prepareIndex("employees").setId("1").setSource("employee_id", 1000000000001L), + client().prepareIndex("employees").setId("2").setSource("employee_id", 2000000000002L), + client().prepareIndex("employees").setId("3").setSource("employee_id", new long[] { 1000000000001L, 3000000000003L }), + client().prepareIndex("employees").setId("4").setSource("employee_id", 4000000000004L) + ); + refresh(); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.addLong(1000000000001L); + bitmap.addLong(4000000000004L); + + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + DataOutputStream dos = new DataOutputStream(bos); + bitmap.serializePortable(dos); + dos.close(); + + BytesArray bitmapBytes = new BytesArray(bos.toByteArray()); + + // directly building the terms query builder, so pass in the bitmap value as BytesArray + SearchResponse searchResponse = client().prepareSearch("employees") + .setQuery(constantScoreQuery(termsQuery("employee_id", bitmapBytes).valueType(TermsQueryBuilder.ValueType.BITMAP))) + .get(); + assertHitCount(searchResponse, 3L); + assertSearchHits(searchResponse, "1", "3", "4"); + } + + public void testTermsQueryWithBitmap64IndexAndDocValues() throws Exception { + assertAcked( + prepareCreate("employees2").setMapping( + jsonBuilder().startObject() + .startObject("properties") + .startObject("employee_id") + .field("type", "long") + // Both index and doc values enabled (default) + .endObject() + .endObject() + .endObject() + ) + ); + + indexRandom( + true, + client().prepareIndex("employees2").setId("1").setSource("employee_id", 1000000000001L), + client().prepareIndex("employees2").setId("2").setSource("employee_id", 2000000000002L), + client().prepareIndex("employees2").setId("3").setSource("employee_id", 3000000000003L) + ); + refresh(); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.addLong(1000000000001L); + bitmap.addLong(3000000000003L); + + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + try (DataOutputStream dos = new DataOutputStream(bos)) { + bitmap.serializePortable(dos); + } + + BytesArray bitmapBytes = new BytesArray(bos.toByteArray()); + + SearchResponse searchResponse = client().prepareSearch("employees2") + .setQuery(constantScoreQuery(termsQuery("employee_id", bitmapBytes).valueType(TermsQueryBuilder.ValueType.BITMAP))) + .get(); + + assertHitCount(searchResponse, 2L); + assertSearchHits(searchResponse, "1", "3"); + } + public void testTermsLookupFilter() throws Exception { assertAcked(prepareCreate("lookup").setMapping("terms", "type=text", "other", "type=text")); indexRandomForConcurrentSearch("lookup"); 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 321ae22158bc9..686885b67eb63 100644 --- a/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java @@ -74,9 +74,13 @@ import org.opensearch.search.approximate.ApproximatePointRangeQuery; import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.search.lookup.SearchLookup; +import org.opensearch.search.query.Bitmap64DocValuesQuery; +import org.opensearch.search.query.Bitmap64IndexQuery; import org.opensearch.search.query.BitmapDocValuesQuery; import org.opensearch.search.query.BitmapIndexQuery; +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; import java.io.IOException; import java.math.BigInteger; import java.nio.ByteBuffer; @@ -93,6 +97,7 @@ import java.util.function.Supplier; import org.roaringbitmap.RoaringBitmap; +import org.roaringbitmap.longlong.Roaring64NavigableMap; /** * A {@link FieldMapper} for numeric types: byte, short, int, long, float, double and unsigned long. @@ -1368,6 +1373,29 @@ public Query termsQuery(String field, List values, boolean hasDocValues, return LongPoint.newSetQuery(field, v); } + @Override + public Query bitmapQuery(String field, BytesArray bitmapArray, boolean isSearchable, boolean hasDocValues) { + // Extract bytes safely + BytesRef ref = bitmapArray.toBytesRef(); + byte[] bytes = Arrays.copyOfRange(ref.bytes, ref.offset, ref.offset + ref.length); + + Roaring64NavigableMap bitmap64 = new Roaring64NavigableMap(); + try { + bitmap64.deserializePortable(new DataInputStream(new ByteArrayInputStream(bytes))); + } catch (IOException e) { + throw new IllegalArgumentException("Failed to deserialize the 64-bit bitmap.", e); + } + + // Note: bitmap64 instance is safely shared between queries as both perform read-only operations + if (isSearchable && hasDocValues) { + return new IndexOrDocValuesQuery(new Bitmap64IndexQuery(field, bitmap64), new Bitmap64DocValuesQuery(field, bitmap64)); + } + if (isSearchable) { + return new Bitmap64IndexQuery(field, bitmap64); + } + return new Bitmap64DocValuesQuery(field, bitmap64); + } + @Override public Query rangeQuery( String field, diff --git a/server/src/main/java/org/opensearch/search/query/Bitmap64DocValuesQuery.java b/server/src/main/java/org/opensearch/search/query/Bitmap64DocValuesQuery.java new file mode 100644 index 0000000000000..13d48897ae348 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/query/Bitmap64DocValuesQuery.java @@ -0,0 +1,158 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.query; + +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.ConstantScoreScorer; +import org.apache.lucene.search.ConstantScoreWeight; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchNoDocsQuery; +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.TwoPhaseIterator; +import org.apache.lucene.search.Weight; +import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.RamUsageEstimator; + +import java.io.IOException; +import java.util.Objects; + +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +import static org.opensearch.search.query.Bitmap64IndexQuery.checkArgs; + +/** + * 64-bit Bitmap DocValues Query + * Same logic as BitmapDocValuesQuery but supports long values. + */ +public class Bitmap64DocValuesQuery extends Query implements Accountable { + + final String field; + final Roaring64NavigableMap bitmap; + final long min; + final long max; + + public Bitmap64DocValuesQuery(String field, Roaring64NavigableMap bitmap) { + checkArgs(field, bitmap); + this.field = field; + this.bitmap = bitmap; + if (!bitmap.isEmpty()) { + min = bitmap.first(); + max = bitmap.last(); + } else { + min = 0; + max = 0; + } + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + return new ConstantScoreWeight(this, boost) { + @Override + public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException { + SortedNumericDocValues values = DocValues.getSortedNumeric(context.reader(), field); + 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 value >= min && value <= max && bitmap.contains(value); + } + + @Override + public float matchCost() { + return 5; + } + }; + } 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 (value < min) { + continue; + } else if (value > max) { + return false; + } else if (bitmap.contains(value)) { + return true; + } + } + return false; + } + + @Override + public float matchCost() { + return 5; + } + }; + } + + final Scorer scorer = new ConstantScoreScorer(score(), scoreMode, iterator); + return new Weight.DefaultScorerSupplier(scorer); + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return DocValues.isCacheable(ctx, field); + } + }; + } + + @Override + public String toString(String field) { + return "Bitmap64DocValuesQuery(field=" + this.field + ")"; + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + if (bitmap.isEmpty()) { + return new MatchNoDocsQuery(); + } + return super.rewrite(indexSearcher); + } + + @Override + public boolean equals(Object other) { + if (sameClassAs(other) == false) { + return false; + } + Bitmap64DocValuesQuery that = (Bitmap64DocValuesQuery) other; + return field.equals(that.field) && bitmap.equals(that.bitmap); + } + + @Override + public int hashCode() { + return Objects.hash(classHash(), field, bitmap); + } + + @Override + public long ramBytesUsed() { + return RamUsageEstimator.shallowSizeOfInstance(Bitmap64DocValuesQuery.class) + RamUsageEstimator.sizeOf(field) + bitmap + .getLongSizeInBytes(); + } + + @Override + public void visit(QueryVisitor visitor) { + if (visitor.acceptField(field)) { + visitor.visitLeaf(this); + } + } +} diff --git a/server/src/main/java/org/opensearch/search/query/Bitmap64IndexQuery.java b/server/src/main/java/org/opensearch/search/query/Bitmap64IndexQuery.java new file mode 100644 index 0000000000000..a06195da5a714 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/query/Bitmap64IndexQuery.java @@ -0,0 +1,246 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.query; + +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.PointValues; +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.MatchNoDocsQuery; +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.Accountable; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.BytesRefIterator; +import org.apache.lucene.util.DocIdSetBuilder; +import org.apache.lucene.util.RamUsageEstimator; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; + +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +/** + * A query that matches all documents that contain a set of long values represented by a 64-bit bitmap + * + * @opensearch.internal + */ +public class Bitmap64IndexQuery extends Query implements Accountable { + + private final Roaring64NavigableMap bitmap; + private final String field; + + public Bitmap64IndexQuery(String field, Roaring64NavigableMap bitmap) { + checkArgs(field, bitmap); + this.bitmap = bitmap; + this.field = field; + } + + static void checkArgs(String field, Roaring64NavigableMap bitmap) { + if (field == null) throw new IllegalArgumentException("field must not be null"); + if (bitmap == null) throw new IllegalArgumentException("bitmap must not be null"); + } + + interface BitmapIterator extends BytesRefIterator { + BytesRef next(); + + void advance(byte[] target); + } + + private static BitmapIterator bitmapEncodedIterator(Roaring64NavigableMap bitmap) { + return new BitmapIterator() { + private final org.roaringbitmap.longlong.LongIterator it = bitmap.getLongIterator(); + private final BytesRef encoded = new BytesRef(new byte[Long.BYTES]); + private final byte[] currentBytes = new byte[Long.BYTES]; + private boolean hasBuffered = false; + + @Override + public BytesRef next() { + if (hasBuffered) { + hasBuffered = false; + System.arraycopy(currentBytes, 0, encoded.bytes, 0, Long.BYTES); + return encoded; + } + + if (!it.hasNext()) return null; + + long v = it.next(); + LongPoint.encodeDimension(v, encoded.bytes, 0); + return encoded; + } + + @Override + public void advance(byte[] target) { + while (it.hasNext()) { + long v = it.next(); + LongPoint.encodeDimension(v, currentBytes, 0); + + if (Arrays.compareUnsigned(currentBytes, target) >= 0) { + hasBuffered = true; + return; + } + } + } + }; + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) { + return new ConstantScoreWeight(this, boost) { + + final long cardinality = bitmap.getLongCardinality(); + + @Override + public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException { + LeafReader reader = context.reader(); + PointValues values = reader.getPointValues(field); + if (values == null) return null; + + if (values.getNumIndexDimensions() != 1) { + throw new IllegalArgumentException("field must have only one dimension"); + } + + return new ScorerSupplier() { + long cost = -1; + final DocIdSetBuilder result = new DocIdSetBuilder(reader.maxDoc(), values); + final MergePointVisitor visitor = new MergePointVisitor(result); + + @Override + public Scorer get(long leadCost) throws IOException { + values.intersect(visitor); + return new ConstantScoreScorer(score(), scoreMode, result.build().iterator()); + } + + @Override + public long cost() { + if (cost == -1) cost = cardinality * 20; // same heuristic as 32-bit + return cost; + } + }; + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return true; // depends only on segment points + } + }; + } + + private class MergePointVisitor implements PointValues.IntersectVisitor { + + private final DocIdSetBuilder result; + private final BitmapIterator iterator; + private BytesRef nextQueryPoint; + private final ArrayUtil.ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(Long.BYTES); + private DocIdSetBuilder.BulkAdder adder; + + MergePointVisitor(DocIdSetBuilder result) { + this.result = result; + this.iterator = bitmapEncodedIterator(bitmap); + nextQueryPoint = iterator.next(); + } + + @Override + public void grow(int count) { + adder = result.grow(count); + } + + @Override + public void visit(int docID) { + adder.add(docID); + } + + @Override + public void visit(DocIdSetIterator iterator) throws IOException { + adder.add(iterator); + } + + @Override + public void visit(int docID, byte[] packedValue) { + if (matches(packedValue)) visit(docID); + } + + @Override + public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException { + if (matches(packedValue)) adder.add(iterator); + } + + private boolean matches(byte[] packedValue) { + while (nextQueryPoint != null) { + int cmp = comparator.compare(nextQueryPoint.bytes, nextQueryPoint.offset, packedValue, 0); + if (cmp == 0) return true; + if (cmp < 0) { + iterator.advance(packedValue); + nextQueryPoint = iterator.next(); + } else break; + } + return false; + } + + @Override + public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { + while (nextQueryPoint != null) { + int cmpMin = comparator.compare(nextQueryPoint.bytes, nextQueryPoint.offset, minPackedValue, 0); + if (cmpMin < 0) { + iterator.advance(minPackedValue); + nextQueryPoint = iterator.next(); + continue; + } + int cmpMax = comparator.compare(nextQueryPoint.bytes, nextQueryPoint.offset, maxPackedValue, 0); + if (cmpMax > 0) return PointValues.Relation.CELL_OUTSIDE_QUERY; + return PointValues.Relation.CELL_CROSSES_QUERY; + } + return PointValues.Relation.CELL_OUTSIDE_QUERY; + } + } + + @Override + public Query rewrite(IndexSearcher searcher) throws IOException { + if (bitmap.isEmpty()) return new MatchNoDocsQuery(); + return super.rewrite(searcher); + } + + @Override + public void visit(QueryVisitor visitor) { + if (visitor.acceptField(field)) visitor.visitLeaf(this); + } + + @Override + public String toString(String field) { + return "Bitmap64IndexQuery(field=" + this.field + ")"; + } + + @Override + public boolean equals(Object other) { + if (sameClassAs(other) == false) return false; + Bitmap64IndexQuery that = (Bitmap64IndexQuery) other; + return field.equals(that.field) && bitmap.equals(that.bitmap); + } + + @Override + public int hashCode() { + return Objects.hash(classHash(), field, bitmap); + } + + @Override + public long ramBytesUsed() { + return RamUsageEstimator.shallowSizeOfInstance(Bitmap64IndexQuery.class) + RamUsageEstimator.sizeOf(field) + bitmap + .getLongSizeInBytes(); + } +} diff --git a/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java index 5f2ba1f083cc5..ced21855038bf 100644 --- a/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java @@ -77,11 +77,15 @@ import org.opensearch.search.MultiValueMode; import org.opensearch.search.approximate.ApproximatePointRangeQuery; import org.opensearch.search.approximate.ApproximateScoreQuery; +import org.opensearch.search.query.Bitmap64DocValuesQuery; +import org.opensearch.search.query.Bitmap64IndexQuery; import org.opensearch.search.query.BitmapDocValuesQuery; import org.opensearch.search.query.BitmapIndexQuery; import org.junit.Before; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; import java.io.IOException; import java.math.BigDecimal; import java.math.BigInteger; @@ -93,6 +97,7 @@ import java.util.function.Supplier; import org.roaringbitmap.RoaringBitmap; +import org.roaringbitmap.longlong.Roaring64NavigableMap; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.either; @@ -1024,6 +1029,48 @@ public void testBitmapQuery() throws IOException { assertThrows(IllegalArgumentException.class, () -> finalFt.bitmapQuery(bitmap)); } + public void testBitmapQuery64() throws IOException { + Roaring64NavigableMap r = new Roaring64NavigableMap(); + byte[] array; + + try (ByteArrayOutputStream bos = new ByteArrayOutputStream(); DataOutputStream dos = new DataOutputStream(bos)) { + r.serializePortable(dos); + dos.flush(); + array = bos.toByteArray(); + } + + BytesArray bitmap = new BytesArray(array); + + NumberFieldType ft = new NumberFieldMapper.NumberFieldType("field", NumberType.LONG); + + assertEquals( + new IndexOrDocValuesQuery(new Bitmap64IndexQuery("field", r), new Bitmap64DocValuesQuery("field", r)), + ft.bitmapQuery(bitmap) + ); + + ft = new NumberFieldType("field", NumberType.LONG, false, false, true, true, true, null, Collections.emptyMap()); + assertEquals(new Bitmap64DocValuesQuery("field", r), ft.bitmapQuery(bitmap)); + + ft = new NumberFieldType("field", NumberType.LONG, true, false, false, false, true, null, Collections.emptyMap()); + assertEquals(new Bitmap64IndexQuery("field", r), ft.bitmapQuery(bitmap)); + + Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, new IndexWriterConfig()); + DirectoryReader reader = DirectoryReader.open(w); + + assertEquals(new MatchNoDocsQuery(), ft.bitmapQuery(bitmap).rewrite(newSearcher(reader))); + + reader.close(); + w.close(); + dir.close(); + + NumberType type = randomValueOtherThan(NumberType.LONG, () -> randomFrom(NumberType.values())); + ft = new NumberFieldMapper.NumberFieldType("field", type); + NumberFieldType finalFt = ft; + + assertThrows(IllegalArgumentException.class, () -> finalFt.bitmapQuery(bitmap)); + } + public void testFetchUnsignedLongDocValues() throws IOException { Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, new IndexWriterConfig(null)); diff --git a/server/src/test/java/org/opensearch/search/query/Bitmap64DocValuesQueryTests.java b/server/src/test/java/org/opensearch/search/query/Bitmap64DocValuesQueryTests.java new file mode 100644 index 0000000000000..aaf2256dce0f2 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/query/Bitmap64DocValuesQueryTests.java @@ -0,0 +1,211 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.query; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchNoDocsQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Set; + +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +/** + * Tests for {@link Bitmap64DocValuesQuery} + */ +public class Bitmap64DocValuesQueryTests extends OpenSearchTestCase { + + /** Single value per doc */ + public void testSingleValuePerDoc() throws Exception { + Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig()); + + addDoc(writer, 1L); + addDoc(writer, 2L); + addDoc(writer, 3L); + addDoc(writer, 4L); + + writer.close(); + + IndexReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = newSearcher(reader); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(1L); + bitmap.add(4L); + + Bitmap64DocValuesQuery query = new Bitmap64DocValuesQuery("product_id", bitmap); + TopDocs topDocs = searcher.search(query, 10); + + assertEquals(2, topDocs.totalHits.value()); + + Set matchedDocs = new HashSet<>(); + for (int i = 0; i < topDocs.scoreDocs.length; i++) { + matchedDocs.add(topDocs.scoreDocs[i].doc); + } + assertTrue(matchedDocs.contains(0)); + assertTrue(matchedDocs.contains(3)); + + reader.close(); + dir.close(); + } + + /** Multi-value per doc */ + public void testMultiValuePerDoc() throws Exception { + Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig()); + + addDoc(writer, 1L); + + Document doc = new Document(); + doc.add(new LongPoint("product_id", 2L)); + doc.add(new SortedNumericDocValuesField("product_id", 2L)); + doc.add(new LongPoint("product_id", 3L)); + doc.add(new SortedNumericDocValuesField("product_id", 3L)); + writer.addDocument(doc); + + addDoc(writer, 3L); + + addDoc(writer, 4L); + + writer.close(); + + IndexReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = newSearcher(reader); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(3L); + + Bitmap64DocValuesQuery query = new Bitmap64DocValuesQuery("product_id", bitmap); + TopDocs topDocs = searcher.search(query, 10); + + assertEquals(2, topDocs.totalHits.value()); + + Set matchedDocs = new HashSet<>(); + for (int i = 0; i < topDocs.scoreDocs.length; i++) { + matchedDocs.add(topDocs.scoreDocs[i].doc); + } + assertEquals(Set.of(1, 2), matchedDocs); + + reader.close(); + dir.close(); + } + + public void testEmptyBitmap() throws Exception { + Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig()); + + addDoc(writer, 42L); + writer.close(); + + IndexReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = newSearcher(reader); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + Query query = new Bitmap64DocValuesQuery("product_id", bitmap); + + TopDocs topDocs = searcher.search(query, 10); + assertEquals(0, topDocs.totalHits.value()); + + reader.close(); + dir.close(); + } + + public void testEmptyBitmapRewritesToMatchNoDocs() throws Exception { + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + Query query = new Bitmap64DocValuesQuery("product_id", bitmap); + + Query rewritten = query.rewrite(null); + assertTrue(rewritten instanceof MatchNoDocsQuery); + } + + public void testRangeOptimization() throws Exception { + Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig()); + + for (long i = 0; i < 100; i++) { + addDoc(writer, i); + } + writer.close(); + + IndexReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = newSearcher(reader); + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + for (long i = 10; i <= 20; i++) { + bitmap.add(i); + } + + Query query = new Bitmap64DocValuesQuery("product_id", bitmap); + TopDocs topDocs = searcher.search(query, 20); + + assertEquals(11, topDocs.totalHits.value()); + + reader.close(); + dir.close(); + } + + public void testLargeValues() throws Exception { + Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig()); + + long[] values = { Long.MAX_VALUE - 100, Long.MAX_VALUE - 50, Long.MAX_VALUE - 10, Long.MAX_VALUE - 1 }; + + for (long value : values) { + addDoc(writer, value); + } + writer.close(); + + IndexReader reader = DirectoryReader.open(dir); + IndexSearcher searcher = newSearcher(reader); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(Long.MAX_VALUE - 50); + bitmap.add(Long.MAX_VALUE - 10); + + Query query = new Bitmap64DocValuesQuery("product_id", bitmap); + TopDocs topDocs = searcher.search(query, 10); + + assertEquals(2, topDocs.totalHits.value()); + + reader.close(); + dir.close(); + } + + public void testNullFieldThrowsException() { + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(1L); + + expectThrows(IllegalArgumentException.class, () -> { new Bitmap64DocValuesQuery(null, bitmap); }); + } + + public void testNullBitmapThrowsException() { + expectThrows(IllegalArgumentException.class, () -> { new Bitmap64DocValuesQuery("field", null); }); + } + + /** + * Helper method to add a document with both point values and doc values + */ + private void addDoc(IndexWriter writer, long value) throws IOException { + Document doc = new Document(); + doc.add(new LongPoint("product_id", value)); + doc.add(new SortedNumericDocValuesField("product_id", value)); + writer.addDocument(doc); + } +} diff --git a/server/src/test/java/org/opensearch/search/query/Bitmap64IndexQueryTests.java b/server/src/test/java/org/opensearch/search/query/Bitmap64IndexQueryTests.java new file mode 100644 index 0000000000000..8d0742b28d2b1 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/query/Bitmap64IndexQueryTests.java @@ -0,0 +1,162 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.query; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.Weight; +import org.apache.lucene.store.Directory; +import org.opensearch.common.Randomness; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; + +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +public class Bitmap64IndexQueryTests extends OpenSearchTestCase { + + private Directory dir; + private IndexWriter w; + private DirectoryReader reader; + private IndexSearcher searcher; + + @Before + public void initSearcher() throws IOException { + dir = newDirectory(); + w = new IndexWriter(dir, newIndexWriterConfig()); + reader = DirectoryReader.open(w); + } + + @After + public void closeAll() throws IOException { + reader.close(); + w.close(); + dir.close(); + } + + public void testScore() throws IOException { + addDoc(1L); + addDoc(2L); + addDoc(3L); + addDoc(4L); + + refresh(); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(1L); + bitmap.add(4L); + + Bitmap64IndexQuery query = new Bitmap64IndexQuery("product_id", bitmap); + Weight weight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f); + + List actual = getMatchingValues(weight, reader); + assertEquals(List.of(1L, 4L), actual); + } + + public void testScoreMultiValues() throws IOException { + addDoc(1L); + addDoc(2L, 3L); + addDoc(3L); + addDoc(4L); + + refresh(); + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + bitmap.add(3L); + + Bitmap64IndexQuery query = new Bitmap64IndexQuery("product_id", bitmap); + Weight weight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f); + + Set actual = new HashSet<>(getMatchingValues(weight, reader)); + assertEquals(Set.of(2L, 3L), actual); + } + + public void testRandomDocumentsAndQueries() throws IOException { + Random random = Randomness.get(); + int valueRange = 10_000; + + for (long i = 0; i <= valueRange; i++) { + addDoc(i); + } + + refresh(); + + Set queryValues = new HashSet<>(); + for (int i = 0; i < 5; i++) { + queryValues.add((long) random.nextInt(valueRange)); + } + + Roaring64NavigableMap bitmap = new Roaring64NavigableMap(); + queryValues.forEach(bitmap::add); + + Bitmap64IndexQuery query = new Bitmap64IndexQuery("product_id", bitmap); + Weight weight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f); + + Set actual = new HashSet<>(getMatchingValues(weight, reader)); + assertEquals(queryValues, actual); + } + + // ---------------- Helpers ---------------- + + private void addDoc(long... values) throws IOException { + Document d = new Document(); + for (long v : values) { + d.add(new LongPoint("product_id", v)); + d.add(new SortedNumericDocValuesField("product_id", v)); + } + w.addDocument(d); + } + + private void refresh() throws IOException { + w.commit(); + reader.close(); + reader = DirectoryReader.open(w); + searcher = newSearcher(reader); + } + + static List getMatchingValues(Weight weight, IndexReader reader) throws IOException { + List actual = new ArrayList<>(); + for (LeafReaderContext leaf : reader.leaves()) { + SortedNumericDocValues dv = DocValues.getSortedNumeric(leaf.reader(), "product_id"); + Scorer scorer = weight.scorer(leaf); + if (scorer == null) continue; + + DocIdSetIterator it = scorer.iterator(); + int docId; + while ((docId = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) { + if (dv.advanceExact(docId)) { + for (int i = 0; i < dv.docValueCount(); i++) { + actual.add(dv.nextValue()); + } + } + } + } + Collections.sort(actual); + return actual; + } +}