diff --git a/CHANGELOG.md b/CHANGELOG.md index 56b1e7b4ffba3..141278480fb1a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add _list/shards API as paginated alternate to _cat/shards ([#14641](https://github.com/opensearch-project/OpenSearch/pull/14641)) - Latency and Memory allocation improvements to Multi Term Aggregation queries ([#14993](https://github.com/opensearch-project/OpenSearch/pull/14993)) - Flat object field use IndexOrDocValuesQuery to optimize query ([#14383](https://github.com/opensearch-project/OpenSearch/issues/14383)) +- [Experimental] Approximate match_all query with sort by range query ([#16321](https://github.com/opensearch-project/OpenSearch/pull/16321)) ### Dependencies - Bump `com.azure:azure-identity` from 1.13.0 to 1.13.2 ([#15578](https://github.com/opensearch-project/OpenSearch/pull/15578)) @@ -85,7 +86,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Streaming Indexing] Fix intermittent 'The bulk request must be terminated by a newline [\n]' failures [#16337](https://github.com/opensearch-project/OpenSearch/pull/16337)) - Fix wrong default value when setting `index.number_of_routing_shards` to null on index creation ([#16331](https://github.com/opensearch-project/OpenSearch/pull/16331)) - Fix disk usage exceeds threshold cluster can't spin up issue ([#15258](https://github.com/opensearch-project/OpenSearch/pull/15258))) - - Fix inefficient Stream API call chains ending with count() ([#15386](https://github.com/opensearch-project/OpenSearch/pull/15386)) ### Security diff --git a/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java b/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java index 59d999798868e..f14872ef9d9d1 100644 --- a/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java @@ -39,6 +39,7 @@ protected FeatureFlagSettings( FeatureFlags.STAR_TREE_INDEX_SETTING, FeatureFlags.APPLICATION_BASED_CONFIGURATION_TEMPLATES_SETTING, FeatureFlags.READER_WRITER_SPLIT_EXPERIMENTAL_SETTING, - FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING + FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING, + FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING // TODO: Copy set from FeatureFlags.ALL_FEATURE_FLAG_SETTINGS ); } diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index 6df68013a8119..6c941a2417fc7 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -138,7 +138,8 @@ public class FeatureFlags { STAR_TREE_INDEX_SETTING, APPLICATION_BASED_CONFIGURATION_TEMPLATES_SETTING, READER_WRITER_SPLIT_EXPERIMENTAL_SETTING, - TERM_VERSION_PRECOMMIT_ENABLE_SETTING + TERM_VERSION_PRECOMMIT_ENABLE_SETTING, + APPROXIMATE_POINT_RANGE_QUERY_SETTING ); /** diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java index 7fbb38c47572c..9ab27d4ef96a9 100644 --- a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java @@ -492,13 +492,9 @@ public Query rangeQuery( name(), pack(new long[] { l }).bytes, pack(new long[] { u }).bytes, - new long[] { l }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { l }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ); } return query; 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 43e975f95757b..3e1402ee0ea9a 100644 --- a/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java @@ -60,6 +60,7 @@ import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Setting.Property; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.core.xcontent.XContentParser.Token; @@ -71,6 +72,8 @@ import org.opensearch.index.fielddata.plain.SortedNumericIndexFieldData; import org.opensearch.index.query.QueryShardContext; import org.opensearch.search.DocValueFormat; +import org.opensearch.search.approximate.ApproximatePointRangeQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.search.lookup.SearchLookup; import org.opensearch.search.query.BitmapDocValuesQuery; @@ -1064,24 +1067,36 @@ public Query rangeQuery( QueryShardContext context ) { return longRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, (l, u) -> { - if (isSearchable && hasDocValues) { - Query query = LongPoint.newRangeQuery(field, l, u); - Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u); - query = new IndexOrDocValuesQuery(query, dvQuery); - if (context.indexSortedOnField(field)) { - query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query); + Query dvQuery = hasDocValues ? SortedNumericDocValuesField.newSlowRangeQuery(field, l, u) : null; + if (isSearchable) { + Query pointRangeQuery = LongPoint.newRangeQuery(field, l, u); + Query query; + if (dvQuery != null) { + query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery); + if (context.indexSortedOnField(field)) { + query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query); + } + } else { + query = pointRangeQuery; } - return query; - } - if (hasDocValues) { - Query query = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u); - if (context.indexSortedOnField(field)) { - query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query); + if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) { + return new ApproximateScoreQuery( + query, + new ApproximatePointRangeQuery( + field, + LongPoint.pack(new long[] { l }).bytes, + LongPoint.pack(new long[] { u }).bytes, + new long[] { l }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) + ); } return query; } - return LongPoint.newRangeQuery(field, l, u); - + if (context.indexSortedOnField(field)) { + dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, dvQuery); + } + return dvQuery; }); } diff --git a/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java index c62ee0ac39584..bfb7c1696304e 100644 --- a/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java @@ -34,12 +34,15 @@ import org.apache.lucene.search.Query; import org.opensearch.common.lucene.search.Queries; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.core.common.ParsingException; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.ObjectParser; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import java.io.IOException; @@ -88,7 +91,11 @@ public static MatchAllQueryBuilder fromXContent(XContentParser parser) { @Override protected Query doToQuery(QueryShardContext context) { - return Queries.newMatchAllQuery(); + Query query = Queries.newMatchAllQuery(); + if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) { + return new ApproximateScoreQuery(query, new ApproximateMatchAllQuery()); + } + return query; } @Override diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java new file mode 100644 index 0000000000000..ad5908a56a600 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -0,0 +1,90 @@ +/* + * 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.approximate; + +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.sort.FieldSortBuilder; + +import java.io.IOException; +import java.util.Objects; + +/** + * Replaces match-all query with a less expensive query if possible. + *

+ * Currently, will rewrite to a bounded range query over the high/low end of a field if a primary sort is specified + * on that field. + */ +public class ApproximateMatchAllQuery extends ApproximateQuery { + private ApproximateQuery approximation = null; + + @Override + protected boolean canApproximate(SearchContext context) { + approximation = null; + if (context == null) { + return false; + } + if (context.aggregations() != null) { + return false; + } + + if (context.request() != null && context.request().source() != null) { + FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); + if (primarySortField != null && primarySortField.missing() == null) { + MappedFieldType mappedFieldType = context.getQueryShardContext().fieldMapper(primarySortField.fieldName()); + Query rangeQuery = mappedFieldType.rangeQuery(null, null, false, false, null, null, null, context.getQueryShardContext()); + if (rangeQuery instanceof ApproximateScoreQuery) { + ApproximateScoreQuery approximateScoreQuery = (ApproximateScoreQuery) rangeQuery; + approximateScoreQuery.setContext(context); + if (approximateScoreQuery.resolvedQuery instanceof ApproximateQuery) { + approximation = (ApproximateQuery) approximateScoreQuery.resolvedQuery; + return true; + } + } + } + } + return false; + } + + @Override + public String toString(String field) { + return "Approximate(*:*)"; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.visitLeaf(this); + + } + + @Override + public boolean equals(Object o) { + if (sameClassAs(o)) { + ApproximateMatchAllQuery other = (ApproximateMatchAllQuery) o; + return Objects.equals(approximation, other.approximation); + } + return false; + } + + @Override + public int hashCode() { + return classHash(); + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + if (approximation == null) { + throw new IllegalStateException("rewrite called without setting context or query could not be approximated"); + } + return approximation.rewrite(indexSearcher); + } +} diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java index 6ff01f5f39d36..2c1dbf09b9e35 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java @@ -8,6 +8,7 @@ package org.opensearch.search.approximate; +import org.apache.lucene.document.LongPoint; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; @@ -24,41 +25,57 @@ import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.DocIdSetBuilder; import org.apache.lucene.util.IntsRef; -import org.opensearch.index.query.RangeQueryBuilder; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.SortOrder; import java.io.IOException; -import java.util.Arrays; import java.util.Objects; +import java.util.function.Function; /** * An approximate-able version of {@link PointRangeQuery}. It creates an instance of {@link PointRangeQuery} but short-circuits the intersect logic * after {@code size} is hit */ -public abstract class ApproximatePointRangeQuery extends ApproximateQuery { +public class ApproximatePointRangeQuery extends ApproximateQuery { + public static final Function LONG_FORMAT = new Function() { + @Override + public String apply(byte[] bytes) { + return Long.toString(LongPoint.decodeDimension(bytes, 0)); + } + }; + private int size; private SortOrder sortOrder; - public final PointRangeQuery pointRangeQuery; - - protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims) { - this(field, lowerPoint, upperPoint, numDims, 10_000, null); - } + private final PointRangeQuery pointRangeQuery; - protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size) { - this(field, lowerPoint, upperPoint, numDims, size, null); + public ApproximatePointRangeQuery( + String field, + byte[] lowerPoint, + byte[] upperPoint, + int numDims, + Function valueToString + ) { + this(field, lowerPoint, upperPoint, numDims, SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO, null, valueToString); } - protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size, SortOrder sortOrder) { + protected ApproximatePointRangeQuery( + String field, + byte[] lowerPoint, + byte[] upperPoint, + int numDims, + int size, + SortOrder sortOrder, + Function valueToString + ) { this.size = size; this.sortOrder = sortOrder; this.pointRangeQuery = new PointRangeQuery(field, lowerPoint, upperPoint, numDims) { @Override protected String toString(int dimension, byte[] value) { - return super.toString(field); + return valueToString.apply(value); } }; } @@ -435,17 +452,27 @@ public boolean canApproximate(SearchContext context) { } // size 0 could be set for caching if (context.from() + context.size() == 0) { - this.setSize(10_000); + this.setSize(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO); + } else { + this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo())); } - this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo())); if (context.request() != null && context.request().source() != null) { FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); - if (primarySortField != null - && primarySortField.missing() == null - && primarySortField.getFieldName().equals(((RangeQueryBuilder) context.request().source().query()).fieldName())) { - if (primarySortField.order() == SortOrder.DESC) { - this.setSortOrder(SortOrder.DESC); + if (primarySortField != null) { + if (!primarySortField.fieldName().equals(pointRangeQuery.getField())) { + // Cannot sort on a different field. + return false; + } + if (primarySortField.missing() != null) { + // Cannot sort documents missing this field. + return false; } + if (context.request().source().searchAfter() != null) { + // TODO: We *could* optimize searchAfter, especially when this is the only sort field, but existing + // pruning is pretty good. + return false; + } + this.setSortOrder(primarySortField.order()); } } return true; @@ -462,56 +489,16 @@ public final boolean equals(Object o) { } private boolean equalsTo(ApproximatePointRangeQuery other) { - return Objects.equals(pointRangeQuery.getField(), other.pointRangeQuery.getField()) - && pointRangeQuery.getNumDims() == other.pointRangeQuery.getNumDims() - && pointRangeQuery.getBytesPerDim() == other.pointRangeQuery.getBytesPerDim() - && Arrays.equals(pointRangeQuery.getLowerPoint(), other.pointRangeQuery.getLowerPoint()) - && Arrays.equals(pointRangeQuery.getUpperPoint(), other.pointRangeQuery.getUpperPoint()); + return Objects.equals(pointRangeQuery, other.pointRangeQuery); } @Override public final String toString(String field) { final StringBuilder sb = new StringBuilder(); - if (pointRangeQuery.getField().equals(field) == false) { - sb.append(pointRangeQuery.getField()); - sb.append(':'); - } - - // print ourselves as "range per dimension" - for (int i = 0; i < pointRangeQuery.getNumDims(); i++) { - if (i > 0) { - sb.append(','); - } - - int startOffset = pointRangeQuery.getBytesPerDim() * i; - - sb.append('['); - sb.append( - toString( - i, - ArrayUtil.copyOfSubArray(pointRangeQuery.getLowerPoint(), startOffset, startOffset + pointRangeQuery.getBytesPerDim()) - ) - ); - sb.append(" TO "); - sb.append( - toString( - i, - ArrayUtil.copyOfSubArray(pointRangeQuery.getUpperPoint(), startOffset, startOffset + pointRangeQuery.getBytesPerDim()) - ) - ); - sb.append(']'); - } + sb.append("Approximate("); + sb.append(pointRangeQuery.toString()); + sb.append(")"); return sb.toString(); } - - /** - * Returns a string of a single value in a human-readable format for debugging. This is used by - * {@link #toString()}. - * - * @param dimension dimension of the particular value - * @param value single value, never null - * @return human readable value for debugging - */ - protected abstract String toString(int dimension, byte[] value); } diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java index 2395142c606ae..6b39606620716 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java @@ -42,9 +42,10 @@ public ApproximateQuery getApproximationQuery() { } @Override - public final Query rewrite(IndexSearcher indexSearcher) throws IOException { + public Query rewrite(IndexSearcher indexSearcher) throws IOException { if (resolvedQuery == null) { - throw new IllegalStateException("Cannot rewrite resolved query without setContext being called"); + // Default to the original query. This suggests that we were not called from ContextIndexSearcher. + return originalQuery.rewrite(indexSearcher); } return resolvedQuery.rewrite(indexSearcher); } diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java index aa8212e8dad69..813aef218cc34 100644 --- a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java @@ -189,6 +189,9 @@ public void setAggregatedDfs(AggregatedDfs aggregatedDfs) { @Override public Query rewrite(Query original) throws IOException { + if (original instanceof ApproximateScoreQuery) { + ((ApproximateScoreQuery) original).setContext(this.searchContext); + } if (profiler != null) { profiler.startRewriteTime(); } @@ -219,9 +222,6 @@ public Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws profiler.pollLastElement(); } return new ProfileWeight(query, weight, profile); - } else if (query instanceof ApproximateScoreQuery) { - ((ApproximateScoreQuery) query).setContext(searchContext); - return super.createWeight(query, scoreMode, boost); } else { return super.createWeight(query, scoreMode, boost); } diff --git a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java index 15b16f4610062..70922cc06945f 100644 --- a/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java @@ -68,7 +68,10 @@ import org.opensearch.index.query.QueryShardContext; import org.opensearch.search.approximate.ApproximatePointRangeQuery; import org.opensearch.search.approximate.ApproximateScoreQuery; +import org.opensearch.test.FeatureFlagSetter; +import org.opensearch.test.TestSearchContext; import org.joda.time.DateTimeZone; +import org.junit.Before; import java.io.IOException; import java.time.ZoneOffset; @@ -82,6 +85,13 @@ public class DateFieldTypeTests extends FieldTypeTestCase { private static final long nowInMillis = 0; + @Override + @Before + public void setUp() throws Exception { + super.setUp(); + FeatureFlagSetter.set(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY); + } + public void testIsFieldWithinRangeEmptyReader() throws IOException { QueryRewriteContext context = new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis); IndexReader reader = new MultiReader(); @@ -222,13 +232,9 @@ public void testTermQuery() { "field", pack(new long[] { instant }).bytes, pack(new long[] { instant + 999 }).bytes, - new long[] { instant }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { instant }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ); assumeThat( "Using Approximate Range Query as default", @@ -281,32 +287,22 @@ public void testRangeQuery() throws IOException { String date2 = "2016-04-28T11:33:52"; long instant1 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date1)).toInstant().toEpochMilli(); long instant2 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date2)).toInstant().toEpochMilli() + 999; - Query expected = new ApproximateScoreQuery( - new IndexOrDocValuesQuery( - LongPoint.newRangeQuery("field", instant1, instant2), - SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2) - ), - new ApproximatePointRangeQuery( - "field", - pack(new long[] { instant1 }).bytes, - pack(new long[] { instant2 }).bytes, - new long[] { instant1 }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + Query expected = new ApproximatePointRangeQuery( + "field", + pack(new long[] { instant1 }).bytes, + pack(new long[] { instant2 }).bytes, + new long[] { instant1 }.length, + ApproximatePointRangeQuery.LONG_FORMAT ); assumeThat( "Using Approximate Range Query as default", FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY), is(true) ); - assertEquals( - expected, - ft.rangeQuery(date1, date2, true, true, null, null, null, context).rewrite(new IndexSearcher(new MultiReader())) - ); + Query rangeQuery = ft.rangeQuery(date1, date2, true, true, null, null, null, context); + assertTrue(rangeQuery instanceof ApproximateScoreQuery); + ((ApproximateScoreQuery) rangeQuery).setContext(new TestSearchContext(context)); + assertEquals(expected, rangeQuery.rewrite(new IndexSearcher(new MultiReader()))); instant1 = nowInMillis; instant2 = instant1 + 100; @@ -320,13 +316,9 @@ protected String toString(int dimension, byte[] value) { "field", pack(new long[] { instant1 }).bytes, pack(new long[] { instant2 }).bytes, - new long[] { instant1 }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { instant1 }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ) ); assumeThat( @@ -391,23 +383,19 @@ public void testRangeQueryWithIndexSort() { long instant2 = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date2)).toInstant().toEpochMilli() + 999; Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2); - Query expected = new IndexSortSortedNumericDocValuesRangeQuery( - "field", - instant1, - instant2, - new ApproximateScoreQuery( - new IndexOrDocValuesQuery(LongPoint.newRangeQuery("field", instant1, instant2), dvQuery), - new ApproximatePointRangeQuery( - "field", - pack(new long[] { instant1 }).bytes, - pack(new long[] { instant2 }).bytes, - new long[] { instant1 }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + Query expected = new ApproximateScoreQuery( + new IndexSortSortedNumericDocValuesRangeQuery( + "field", + instant1, + instant2, + new IndexOrDocValuesQuery(LongPoint.newRangeQuery("field", instant1, instant2), dvQuery) + ), + new ApproximatePointRangeQuery( + "field", + pack(new long[] { instant1 }).bytes, + pack(new long[] { instant2 }).bytes, + new long[] { instant1 }.length, + ApproximatePointRangeQuery.LONG_FORMAT ) ); assumeThat( diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java index a11a16f421783..c8413df7128a1 100644 --- a/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/RangeFieldQueryStringQueryBuilderTests.java @@ -204,13 +204,9 @@ public void testDateRangeQuery() throws Exception { DATE_FIELD_NAME, pack(new long[] { parser.parse(lowerBoundExact, () -> 0).toEpochMilli() }).bytes, pack(new long[] { parser.parse(upperBoundExact, () -> 0).toEpochMilli() }).bytes, - new long[] { parser.parse(lowerBoundExact, () -> 0).toEpochMilli() }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { parser.parse(lowerBoundExact, () -> 0).toEpochMilli() }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ), queryOnDateField ); diff --git a/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java index 6f295100d9e47..0d7ff60bf9db5 100644 --- a/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java @@ -883,13 +883,9 @@ private ApproximateScoreQuery calculateExpectedDateQuery(long lower, long upper) DATE_FIELD_NAME, pack(new long[] { lower }).bytes, pack(new long[] { upper }).bytes, - new long[] { lower }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { lower }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ); } diff --git a/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java index 79f1452db297b..ec9ba57653e8e 100644 --- a/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/RangeQueryBuilderTests.java @@ -59,8 +59,10 @@ import org.opensearch.search.approximate.ApproximateQuery; import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.AbstractQueryTestCase; +import org.opensearch.test.FeatureFlagSetter; import org.joda.time.DateTime; import org.joda.time.chrono.ISOChronology; +import org.junit.Before; import java.io.IOException; import java.time.Instant; @@ -78,6 +80,13 @@ import static org.junit.Assume.assumeThat; public class RangeQueryBuilderTests extends AbstractQueryTestCase { + @Override + @Before + public void setUp() throws Exception { + super.setUp(); + FeatureFlagSetter.set(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY); + } + @Override protected RangeQueryBuilder doCreateTestQueryBuilder() { RangeQueryBuilder query; @@ -259,13 +268,9 @@ protected void doAssertLuceneQuery(RangeQueryBuilder queryBuilder, Query query, DATE_FIELD_NAME, pack(new long[] { minLong }).bytes, pack(new long[] { maxLong }).bytes, - new long[] { minLong }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { minLong }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ), query ); @@ -355,13 +360,9 @@ public void testDateRangeQueryFormat() throws IOException { DATE_FIELD_NAME, pack(new long[] { lower }).bytes, pack(new long[] { upper }).bytes, - new long[] { lower }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { lower }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ), parsedQuery ); @@ -412,13 +413,9 @@ public void testDateRangeBoundaries() throws IOException { DATE_FIELD_NAME, pack(new long[] { lower }).bytes, pack(new long[] { upper }).bytes, - new long[] { lower }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { lower }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ) , @@ -449,13 +446,9 @@ protected String toString(int dimension, byte[] value) { DATE_FIELD_NAME, pack(new long[] { lower }).bytes, pack(new long[] { upper }).bytes, - new long[] { lower }.length - ) { - @Override - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - } + new long[] { lower }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ) ) , diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java new file mode 100644 index 0000000000000..da3669c249363 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java @@ -0,0 +1,115 @@ +/* + * 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.approximate; + +import org.apache.lucene.search.IndexSearcher; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.BigArrays; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.index.query.QueryShardContext; +import org.opensearch.search.aggregations.AggregatorFactories; +import org.opensearch.search.aggregations.SearchContextAggregations; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.internal.ShardSearchRequest; +import org.opensearch.search.sort.FieldSortBuilder; +import org.opensearch.test.FeatureFlagSetter; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.TestSearchContext; + +import java.io.IOException; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ApproximateMatchAllQueryTests extends OpenSearchTestCase { + + public void testCanApproximate() throws IOException { + ApproximateMatchAllQuery approximateMatchAllQuery = new ApproximateMatchAllQuery(); + // Fail on null searchContext + assertFalse(approximateMatchAllQuery.canApproximate(null)); + + ShardSearchRequest[] shardSearchRequest = new ShardSearchRequest[1]; + + MapperService mockMapper = mock(MapperService.class); + String sortfield = "myfield"; + MappedFieldType myFieldType = new NumberFieldMapper.NumberFieldType(sortfield, NumberFieldMapper.NumberType.LONG); + when(mockMapper.fieldType(sortfield)).thenReturn(myFieldType); + + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1) + .build(); + IndexMetadata indexMetadata = new IndexMetadata.Builder("index").settings(settings).build(); + QueryShardContext queryShardContext = new QueryShardContext( + 0, + new IndexSettings(indexMetadata, settings), + BigArrays.NON_RECYCLING_INSTANCE, + null, + null, + mockMapper, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ); + TestSearchContext searchContext = new TestSearchContext(queryShardContext) { + @Override + public ShardSearchRequest request() { + return shardSearchRequest[0]; + } + }; + + // Fail if aggregations are present + searchContext.aggregations(new SearchContextAggregations(new AggregatorFactories.Builder().build(null, null), null)); + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + searchContext.aggregations(null); + + // Fail on missing ShardSearchRequest + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + + // Fail if source is null or empty + shardSearchRequest[0] = new ShardSearchRequest(null, System.currentTimeMillis(), null); + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + + // Fail if source does not have a sort. + SearchSourceBuilder source = new SearchSourceBuilder(); + shardSearchRequest[0].source(source); + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + + // Still can't approximate, because the APPROXIMATE_POINT_RANGE_QUERY feature is not enabled. + source.sort(sortfield); + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + + // Now we can approximate! + FeatureFlagSetter.set(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY); + assertTrue(approximateMatchAllQuery.canApproximate(searchContext)); + assertTrue(approximateMatchAllQuery.rewrite((IndexSearcher) null) instanceof ApproximatePointRangeQuery); + + // But not if the sort field makes a decision about missing data + source.sorts().clear(); + source.sort(new FieldSortBuilder(sortfield).missing("foo")); + assertFalse(approximateMatchAllQuery.canApproximate(searchContext)); + assertThrows(IllegalStateException.class, () -> approximateMatchAllQuery.rewrite((IndexSearcher) null)); + } + +} diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java index 9c022aade5dc6..081040963c86a 100644 --- a/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java +++ b/server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java @@ -26,12 +26,13 @@ import org.opensearch.search.internal.SearchContext; import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.TestSearchContext; import java.io.IOException; import static java.util.Arrays.asList; +import static org.opensearch.search.approximate.ApproximatePointRangeQuery.LONG_FORMAT; import static org.apache.lucene.document.LongPoint.pack; -import static org.mockito.Mockito.mock; public class ApproximatePointRangeQueryTests extends OpenSearchTestCase { @@ -59,11 +60,13 @@ public void testApproximateRangeEqualsActualRange() throws IOException { try { long lower = RandomNumbers.randomLongBetween(random(), -100, 200); long upper = lower + RandomNumbers.randomLongBetween(random(), 0, 100); - Query approximateQuery = new ApproximatePointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + Query approximateQuery = new ApproximatePointRangeQuery( + "point", + pack(lower).bytes, + pack(upper).bytes, + dims, + LONG_FORMAT + ); Query query = LongPoint.newRangeQuery("point", lower, upper); IndexSearcher searcher = new IndexSearcher(reader); TopDocs topDocs = searcher.search(approximateQuery, 10); @@ -99,11 +102,13 @@ public void testApproximateRangeWithDefaultSize() throws IOException { try { long lower = 0; long upper = 1000; - Query approximateQuery = new ApproximatePointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + Query approximateQuery = new ApproximatePointRangeQuery( + "point", + pack(lower).bytes, + pack(upper).bytes, + dims, + LONG_FORMAT + ); IndexSearcher searcher = new IndexSearcher(reader); TopDocs topDocs = searcher.search(approximateQuery, 10); assertEquals(topDocs.totalHits, new TotalHits(1000, TotalHits.Relation.EQUAL_TO)); @@ -137,11 +142,15 @@ public void testApproximateRangeWithSizeUnderDefault() throws IOException { try { long lower = 0; long upper = 45; - Query approximateQuery = new ApproximatePointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims, 10) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + Query approximateQuery = new ApproximatePointRangeQuery( + "point", + pack(lower).bytes, + pack(upper).bytes, + dims, + 10, + null, + LONG_FORMAT + ); IndexSearcher searcher = new IndexSearcher(reader); TopDocs topDocs = searcher.search(approximateQuery, 10); assertEquals(topDocs.totalHits, new TotalHits(10, TotalHits.Relation.EQUAL_TO)); @@ -180,19 +189,16 @@ public void testApproximateRangeWithSizeOverDefault() throws IOException { pack(lower).bytes, pack(upper).bytes, dims, - 11_000 - ) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + 11_000, + null, + LONG_FORMAT + ); IndexSearcher searcher = new IndexSearcher(reader); TopDocs topDocs = searcher.search(approximateQuery, 11000); - assertEquals(topDocs.totalHits, new TotalHits(11000, TotalHits.Relation.EQUAL_TO)); + assertEquals(new TotalHits(11000, TotalHits.Relation.EQUAL_TO), topDocs.totalHits); } catch (IOException e) { throw new RuntimeException(e); } - } } } @@ -220,11 +226,15 @@ public void testApproximateRangeShortCircuit() throws IOException { try { long lower = 0; long upper = 100; - Query approximateQuery = new ApproximatePointRangeQuery("point", pack(lower).bytes, pack(upper).bytes, dims, 10) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + Query approximateQuery = new ApproximatePointRangeQuery( + "point", + pack(lower).bytes, + pack(upper).bytes, + dims, + 10, + null, + LONG_FORMAT + ); Query query = LongPoint.newRangeQuery("point", lower, upper); ; IndexSearcher searcher = new IndexSearcher(reader); @@ -271,12 +281,9 @@ public void testApproximateRangeShortCircuitAscSort() throws IOException { pack(upper).bytes, dims, 10, - SortOrder.ASC - ) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + SortOrder.ASC, + LONG_FORMAT + ); Query query = LongPoint.newRangeQuery("point", lower, upper); ; IndexSearcher searcher = new IndexSearcher(reader); @@ -305,11 +312,7 @@ protected String toString(int dimension, byte[] value) { } public void testSize() { - ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1, LONG_FORMAT); assertEquals(query.getSize(), 10_000); query.setSize(100); @@ -318,11 +321,7 @@ protected String toString(int dimension, byte[] value) { } public void testSortOrder() { - ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1, LONG_FORMAT); assertNull(query.getSortOrder()); query.setSortOrder(SortOrder.ASC); @@ -330,24 +329,18 @@ protected String toString(int dimension, byte[] value) { } public void testCanApproximate() { - ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + ApproximatePointRangeQuery query = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1, LONG_FORMAT); assertFalse(query.canApproximate(null)); - ApproximatePointRangeQuery queryCanApproximate = new ApproximatePointRangeQuery("point", pack(0).bytes, pack(20).bytes, 1) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - - public boolean canApproximate(SearchContext context) { - return true; - } - }; - SearchContext searchContext = mock(SearchContext.class); + ApproximatePointRangeQuery queryCanApproximate = new ApproximatePointRangeQuery( + "point", + pack(0).bytes, + pack(20).bytes, + 1, + LONG_FORMAT + ); + SearchContext searchContext = new TestSearchContext(null); assertTrue(queryCanApproximate.canApproximate(searchContext)); } } diff --git a/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java b/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java index aa45ea6744227..5af49efb30f68 100644 --- a/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java +++ b/server/src/test/java/org/opensearch/search/approximate/ApproximateScoreQueryTests.java @@ -46,12 +46,9 @@ protected String toString(int dimension, byte[] value) { "test-index", pack(new long[] { l }).bytes, pack(new long[] { u }).bytes, - new long[] { l }.length - ) { - protected String toString(int dimension, byte[] value) { - return Long.toString(LongPoint.decodeDimension(value, 0)); - } - }; + new long[] { l }.length, + ApproximatePointRangeQuery.LONG_FORMAT + ); ApproximateScoreQuery query = new ApproximateScoreQuery(originalQuery, approximateQuery); query.resolvedQuery = approximateQuery;