Skip to content

Commit

Permalink
Remove ApproximateIndexOrDocValuesQuery
Browse files Browse the repository at this point in the history
Looking into the query approximation framework, I realized that we don't
actually need ApproximateIndexOrDocValuesQuery. We already have
ApproximateScoreQuery that is able to choose between the approximate
query and the original query.

Assuming the range query is over an indexed field, it is (potentially)
possible to approximate. If there are doc values, then the original
query will be an IndexOrDocValuesQuery. Otherwise, it will just be a
PointRangeQuery.  Either way, we can wrap the original query into a
generic ApproximateScoreQuery along with the approximate query. The
approximation logic doesn't need to be specialized to the
IndexOrDocValue case.

Signed-off-by: Michael Froh <[email protected]>
  • Loading branch information
msfroh committed Oct 10, 2024
1 parent 96082f7 commit dc2c63a
Show file tree
Hide file tree
Showing 12 changed files with 115 additions and 272 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,8 @@
import org.opensearch.index.query.QueryRewriteContext;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximatePointRangeQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.opensearch.search.lookup.SearchLookup;

import java.io.IOException;
Expand Down Expand Up @@ -113,21 +113,6 @@ public static DateFormatter getDefaultDateTimeFormatter() {
: LEGACY_DEFAULT_DATE_TIME_FORMATTER;
}

public static Query getDefaultQuery(Query pointRangeQuery, Query dvQuery, String name, long l, long u) {
return FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)
? new ApproximateIndexOrDocValuesQuery(
pointRangeQuery,
new ApproximatePointRangeQuery(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));
}
},
dvQuery
)
: new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
}

/**
* Resolution of the date time
*
Expand Down Expand Up @@ -488,22 +473,39 @@ public Query rangeQuery(
}
DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser;
return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> {
Query pointRangeQuery = isSearchable() ? LongPoint.newRangeQuery(name(), l, u) : null;
Query dvQuery = hasDocValues() ? SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u) : null;
if (isSearchable() && hasDocValues()) {
Query query = getDefaultQuery(pointRangeQuery, dvQuery, name(), l, u);
if (context.indexSortedOnField(name())) {
query = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, query);
if (isSearchable()) {
Query pointRangeQuery = LongPoint.newRangeQuery(name(), l, u);
Query query;
if (dvQuery != null) {
query = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
} else {
query = pointRangeQuery;
}
if (FeatureFlags.isEnabled(FeatureFlags.APPROXIMATE_POINT_RANGE_QUERY_SETTING)) {
return new ApproximateScoreQuery(
query,
new ApproximatePointRangeQuery(
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));
}
}
);
}
return query;
}
if (hasDocValues()) {
if (context.indexSortedOnField(name())) {
dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
}
return dvQuery;

// Not searchable. Must have doc values.
if (context.indexSortedOnField(name())) {
dvQuery = new IndexSortSortedNumericDocValuesRangeQuery(name(), l, u, dvQuery);
}
return pointRangeQuery;
return dvQuery;
});
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
import org.opensearch.common.lucene.search.function.FunctionScoreQuery;
import org.opensearch.index.mapper.DateFieldMapper;
import org.opensearch.index.query.DateRangeIncludingNowQuery;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.opensearch.search.internal.SearchContext;

import java.io.IOException;
Expand Down Expand Up @@ -55,7 +55,7 @@ private Helper() {}
queryWrappers.put(FunctionScoreQuery.class, q -> ((FunctionScoreQuery) q).getSubQuery());
queryWrappers.put(DateRangeIncludingNowQuery.class, q -> ((DateRangeIncludingNowQuery) q).getQuery());
queryWrappers.put(IndexOrDocValuesQuery.class, q -> ((IndexOrDocValuesQuery) q).getIndexQuery());
queryWrappers.put(ApproximateIndexOrDocValuesQuery.class, q -> ((ApproximateIndexOrDocValuesQuery) q).getOriginalQuery());
queryWrappers.put(ApproximateScoreQuery.class, q -> ((ApproximateScoreQuery) q).getOriginalQuery());
}

/**
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -433,9 +433,6 @@ public boolean canApproximate(SearchContext context) {
if (context.aggregations() != null) {
return false;
}
if (!(context.query() instanceof ApproximateIndexOrDocValuesQuery)) {
return false;
}
// size 0 could be set for caching
if (context.from() + context.size() == 0) {
this.setSize(10_000);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,12 @@
* Entry-point for the approximation framework.
* This class is heavily inspired by {@link org.apache.lucene.search.IndexOrDocValuesQuery}. It acts as a wrapper that consumer two queries, a regular query and an approximate version of the same. By default, it executes the regular query and returns {@link Weight#scorer} for the original query. At run-time, depending on certain constraints, we can re-write the {@code Weight} to use the approximate weight instead.
*/
public class ApproximateScoreQuery extends Query {
public final class ApproximateScoreQuery extends Query {

private final Query originalQuery;
private final ApproximateQuery approximationQuery;

protected Query resolvedQuery;
Query resolvedQuery;

public ApproximateScoreQuery(Query originalQuery, ApproximateQuery approximationQuery) {
this.originalQuery = originalQuery;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
import org.apache.lucene.index.MultiReader;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexOrDocValuesQuery;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.IndexSortSortedNumericDocValuesRangeQuery;
import org.apache.lucene.search.Query;
Expand All @@ -65,8 +66,8 @@
import org.opensearch.index.query.DateRangeIncludingNowQuery;
import org.opensearch.index.query.QueryRewriteContext;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximatePointRangeQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.joda.time.DateTimeZone;

import java.io.IOException;
Expand Down Expand Up @@ -212,8 +213,11 @@ public void testTermQuery() {
MappedFieldType ft = new DateFieldType("field");
String date = "2015-10-12T14:10:55";
long instant = DateFormatters.from(DateFieldMapper.getDefaultDateTimeFormatter().parse(date)).toInstant().toEpochMilli();
Query expected = new ApproximateIndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant, instant + 999),
Query expected = new ApproximateScoreQuery(
new IndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant, instant + 999),
SortedNumericDocValuesField.newSlowRangeQuery("field", instant, instant + 999)
),
new ApproximatePointRangeQuery(
"field",
pack(new long[] { instant }).bytes,
Expand All @@ -224,8 +228,7 @@ public void testTermQuery() {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
SortedNumericDocValuesField.newSlowRangeQuery("field", instant, instant + 999)
}
);
assumeThat(
"Using Approximate Range Query as default",
Expand Down Expand Up @@ -278,8 +281,11 @@ 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 ApproximateIndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant1, instant2),
Query expected = new ApproximateScoreQuery(
new IndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant1, instant2),
SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
),
new ApproximatePointRangeQuery(
"field",
pack(new long[] { instant1 }).bytes,
Expand All @@ -290,8 +296,7 @@ public void testRangeQuery() throws IOException {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
}
);
assumeThat(
"Using Approximate Range Query as default",
Expand All @@ -306,8 +311,11 @@ protected String toString(int dimension, byte[] value) {
instant1 = nowInMillis;
instant2 = instant1 + 100;
expected = new DateRangeIncludingNowQuery(
new ApproximateIndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant1, instant2),
new ApproximateScoreQuery(
new IndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant1, instant2),
SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
),
new ApproximatePointRangeQuery(
"field",
pack(new long[] { instant1 }).bytes,
Expand All @@ -318,8 +326,7 @@ protected String toString(int dimension, byte[] value) {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
}
)
);
assumeThat(
Expand Down Expand Up @@ -388,8 +395,8 @@ public void testRangeQueryWithIndexSort() {
"field",
instant1,
instant2,
new ApproximateIndexOrDocValuesQuery(
LongPoint.newRangeQuery("field", instant1, instant2),
new ApproximateScoreQuery(
new IndexOrDocValuesQuery(LongPoint.newRangeQuery("field", instant1, instant2), dvQuery),
new ApproximatePointRangeQuery(
"field",
pack(new long[] { instant1 }).bytes,
Expand All @@ -400,8 +407,7 @@ public void testRangeQueryWithIndexSort() {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
dvQuery
}
)
);
assumeThat(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@
import org.opensearch.common.util.FeatureFlags;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.query.QueryStringQueryBuilder;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximatePointRangeQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.opensearch.test.AbstractQueryTestCase;

import java.io.IOException;
Expand Down Expand Up @@ -191,11 +191,14 @@ public void testDateRangeQuery() throws Exception {
is(true)
);
assertEquals(
new ApproximateIndexOrDocValuesQuery(
LongPoint.newRangeQuery(
DATE_FIELD_NAME,
parser.parse(lowerBoundExact, () -> 0).toEpochMilli(),
parser.parse(upperBoundExact, () -> 0).toEpochMilli()
new ApproximateScoreQuery(
new IndexOrDocValuesQuery(
LongPoint.newRangeQuery(
DATE_FIELD_NAME,
parser.parse(lowerBoundExact, () -> 0).toEpochMilli(),
parser.parse(upperBoundExact, () -> 0).toEpochMilli()
),
controlDv
),
new ApproximatePointRangeQuery(
DATE_FIELD_NAME,
Expand All @@ -207,8 +210,7 @@ public void testDateRangeQuery() throws Exception {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
controlDv
}
),
queryOnDateField
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@
import org.opensearch.index.mapper.RangeFieldMapper.RangeFieldType;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.query.QueryShardException;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.opensearch.test.IndexSettingsModule;
import org.joda.time.DateTime;
import org.junit.Before;
Expand Down Expand Up @@ -293,7 +293,7 @@ public void testDateRangeQueryUsingMappingFormatLegacy() {
);
assertEquals(
"field:[1465975790000 TO 1466062190999]",
((IndexOrDocValuesQuery) ((ApproximateIndexOrDocValuesQuery) queryOnDateField).getOriginalQuery()).getIndexQuery().toString()
((IndexOrDocValuesQuery) ((ApproximateScoreQuery) queryOnDateField).getOriginalQuery()).getIndexQuery().toString()
);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
import org.apache.lucene.search.TermQuery;
import org.opensearch.core.common.ParsingException;
import org.opensearch.index.search.MatchQuery.ZeroTermsQuery;
import org.opensearch.search.approximate.ApproximateIndexOrDocValuesQuery;
import org.opensearch.search.approximate.ApproximateScoreQuery;
import org.opensearch.test.AbstractQueryTestCase;

import java.io.IOException;
Expand Down Expand Up @@ -131,7 +131,7 @@ protected void doAssertLuceneQuery(MatchPhraseQueryBuilder queryBuilder, Query q
.or(instanceOf(PointRangeQuery.class))
.or(instanceOf(IndexOrDocValuesQuery.class))
.or(instanceOf(MatchNoDocsQuery.class))
.or(instanceOf(ApproximateIndexOrDocValuesQuery.class))
.or(instanceOf(ApproximateScoreQuery.class))
);
}

Expand Down
Loading

0 comments on commit dc2c63a

Please sign in to comment.