From 23d3985f8897309dfbfc04aa1b7ab3b26e92998a Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Wed, 2 Apr 2025 19:02:43 -0700 Subject: [PATCH 1/7] Adding approximate match all query Signed-off-by: Harsha Vamsi Kalluri --- .../index/query/MatchAllQueryBuilder.java | 4 +- .../approximate/ApproximateMatchAllQuery.java | 90 +++++++++++++++ .../ApproximateMatchAllQueryTests.java | 108 ++++++++++++++++++ 3 files changed, 201 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java create mode 100644 server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java 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..d7bc63db4fecb 100644 --- a/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java @@ -40,6 +40,8 @@ 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 +90,7 @@ public static MatchAllQueryBuilder fromXContent(XContentParser parser) { @Override protected Query doToQuery(QueryShardContext context) { - return Queries.newMatchAllQuery(); + return new ApproximateScoreQuery(Queries.newMatchAllQuery(), new ApproximateMatchAllQuery()); } @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/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..db904751ed76f --- /dev/null +++ b/server/src/test/java/org/opensearch/search/approximate/ApproximateMatchAllQueryTests.java @@ -0,0 +1,108 @@ +/* + * 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.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.BigArrays; +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.search.sort.SortOrder; +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)); + + source.sort(sortfield, SortOrder.ASC); + assertTrue(approximateMatchAllQuery.canApproximate(searchContext)); + assertTrue(approximateMatchAllQuery.rewrite(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(null)); + } + +} From 740e9a0509a40e0d34a6cad16059e2ee67e4d2ee Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Mon, 7 Apr 2025 12:09:00 -0700 Subject: [PATCH 2/7] Minor refactor Signed-off-by: Harsha Vamsi Kalluri --- .../search/approximate/ApproximateMatchAllQuery.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java index ad5908a56a600..0fee7ed51e077 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -42,8 +42,7 @@ protected boolean canApproximate(SearchContext context) { 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; + if (rangeQuery instanceof ApproximateScoreQuery approximateScoreQuery) { approximateScoreQuery.setContext(context); if (approximateScoreQuery.resolvedQuery instanceof ApproximateQuery) { approximation = (ApproximateQuery) approximateScoreQuery.resolvedQuery; From 3a9533957e9c2c38acfe83642b16bfa511520d9e Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Thu, 10 Apr 2025 13:09:39 -0700 Subject: [PATCH 3/7] Fix doc field sorting Signed-off-by: Harsha Vamsi Kalluri --- .../search/approximate/ApproximateMatchAllQuery.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java index 0fee7ed51e077..c30f97b201953 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -39,7 +39,9 @@ protected boolean canApproximate(SearchContext context) { if (context.request() != null && context.request().source() != null) { FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); - if (primarySortField != null && primarySortField.missing() == null) { + if (primarySortField != null + && primarySortField.missing() == null + && !primarySortField.fieldName().equals(FieldSortBuilder.DOC_FIELD_NAME)) { 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) { From 21a40875511ee99c286c052b9bb27fffe6755ebf Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Thu, 10 Apr 2025 14:18:44 -0700 Subject: [PATCH 4/7] Fix tests Signed-off-by: Harsha Vamsi Kalluri --- .../test/indices.validate_query/10_basic.yml | 2 +- .../approximate/ApproximateMatchAllQuery.java | 3 +++ .../index/query/BoolQueryBuilderTests.java | 7 ++++++- .../index/query/MatchAllQueryBuilderTests.java | 7 ++++++- .../index/query/NestedQueryBuilderTests.java | 7 ++++++- .../index/query/WrapperQueryBuilderTests.java | 9 ++++++++- .../FunctionScoreQueryBuilderTests.java | 7 ++++++- .../opensearch/index/search/NestedHelperTests.java | 8 +++++++- .../search/sort/GeoDistanceSortBuilderTests.java | 13 +++++++++++-- .../search/sort/ScriptSortBuilderTests.java | 13 +++++++++++-- 10 files changed, 65 insertions(+), 11 deletions(-) diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml index d3d78c3c60a00..a114ed48741cf 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml @@ -66,7 +66,7 @@ setup: - is_true: valid - match: {_shards.failed: 0} - match: {explanations.0.index: 'testing'} - - match: {explanations.0.explanation: '*:*'} + - match: {explanations.0.explanation: 'ApproximateScoreQuery(originalQuery=*:*, approximationQuery=Approximate(*:*))'} --- "Validate body without query element": diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java index c30f97b201953..6a96b512acc48 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -43,6 +43,9 @@ protected boolean canApproximate(SearchContext context) { && primarySortField.missing() == null && !primarySortField.fieldName().equals(FieldSortBuilder.DOC_FIELD_NAME)) { MappedFieldType mappedFieldType = context.getQueryShardContext().fieldMapper(primarySortField.fieldName()); + if (mappedFieldType == null) { + return false; + } Query rangeQuery = mappedFieldType.rangeQuery(null, null, false, false, null, null, null, context.getQueryShardContext()); if (rangeQuery instanceof ApproximateScoreQuery approximateScoreQuery) { approximateScoreQuery.setContext(context); diff --git a/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java index f3de666c52932..9223d82339252 100644 --- a/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java @@ -42,6 +42,8 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParseException; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.AbstractQueryTestCase; import org.hamcrest.Matchers; @@ -208,7 +210,10 @@ public void testMinShouldMatchFilterWithoutShouldClauses() throws Exception { assertThat(innerBooleanQuery.clauses().size(), equalTo(1)); BooleanClause innerBooleanClause = innerBooleanQuery.clauses().get(0); assertThat(innerBooleanClause.occur(), equalTo(BooleanClause.Occur.MUST)); - assertThat(innerBooleanClause.query(), instanceOf(MatchAllDocsQuery.class)); + assertThat(innerBooleanClause.query(), instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) innerBooleanClause.query(); + assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class)); } public void testMinShouldMatchBiggerThanNumberOfShouldClauses() throws Exception { diff --git a/server/src/test/java/org/opensearch/index/query/MatchAllQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/MatchAllQueryBuilderTests.java index 7e51567ac0f84..458c907766f78 100644 --- a/server/src/test/java/org/opensearch/index/query/MatchAllQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/MatchAllQueryBuilderTests.java @@ -34,6 +34,8 @@ import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.AbstractQueryTestCase; import java.io.IOException; @@ -49,7 +51,10 @@ protected MatchAllQueryBuilder doCreateTestQueryBuilder() { @Override protected void doAssertLuceneQuery(MatchAllQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException { - assertThat(query, instanceOf(MatchAllDocsQuery.class)); + assertThat(query, instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) query; + assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class)); } public void testFromJson() throws IOException { diff --git a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java index c371f004295d6..435e529325646 100644 --- a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java @@ -50,6 +50,8 @@ import org.opensearch.index.mapper.MapperService; import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder; import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.search.fetch.subphase.InnerHitsContext; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.sort.FieldSortBuilder; @@ -493,7 +495,10 @@ public void testNestedDepthAllowed() throws Exception { .filter(c -> c.occur() == BooleanClause.Occur.MUST) .findFirst(); assertTrue(childLeg.isPresent()); - assertEquals(new MatchAllDocsQuery(), childLeg.get().query()); + assertThat(childLeg.get().query(), instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) childLeg.get().query(); + assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class)); }; check.accept(createShardContext()); doWithDepth(randomIntBetween(1, 20), check); diff --git a/server/src/test/java/org/opensearch/index/query/WrapperQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/WrapperQueryBuilderTests.java index 4135f6e0ef049..1786517c1aa1d 100644 --- a/server/src/test/java/org/opensearch/index/query/WrapperQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/WrapperQueryBuilderTests.java @@ -41,12 +41,16 @@ import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.AbstractQueryTestCase; import java.io.IOException; import java.io.UncheckedIOException; import java.io.UnsupportedEncodingException; +import static org.hamcrest.CoreMatchers.instanceOf; + public class WrapperQueryBuilderTests extends AbstractQueryTestCase { @Override @@ -171,7 +175,10 @@ public void testRewriteInnerQueryToo() throws IOException { assertEquals(new TermQuery(new Term(TEXT_FIELD_NAME, "bar")), qb.rewrite(shardContext).toQuery(shardContext)); qb = new WrapperQueryBuilder(new BoolQueryBuilder().toString()); - assertEquals(new MatchAllDocsQuery(), qb.rewrite(shardContext).toQuery(shardContext)); + assertThat(qb.rewrite(shardContext).toQuery(shardContext), instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) qb.rewrite(shardContext).toQuery(shardContext); + assertThat(approxQuery.getOriginalQuery(), instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), instanceOf(ApproximateMatchAllQuery.class)); } @Override diff --git a/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java index 8cf7941941bcb..c762fdcde146d 100644 --- a/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java @@ -67,6 +67,8 @@ import org.opensearch.script.Script; import org.opensearch.script.ScriptType; import org.opensearch.search.MultiValueMode; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.AbstractQueryTestCase; import org.opensearch.test.TestGeoShapeFieldMapperPlugin; import org.joda.time.DateTime; @@ -630,7 +632,10 @@ public void testCustomWeightFactorQueryBuilderWithFunctionScoreWithoutQueryGiven Query parsedQuery = parseQuery(functionScoreQuery(weightFactorFunction(1.3f))).toQuery(createShardContext()); assertThat(parsedQuery, instanceOf(FunctionScoreQuery.class)); FunctionScoreQuery functionScoreQuery = (FunctionScoreQuery) parsedQuery; - assertThat(functionScoreQuery.getSubQuery() instanceof MatchAllDocsQuery, equalTo(true)); + assertThat(functionScoreQuery.getSubQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) functionScoreQuery.getSubQuery(); + assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class)); assertThat((double) (functionScoreQuery.getFunctions()[0]).getWeight(), closeTo(1.3, 0.001)); } diff --git a/server/src/test/java/org/opensearch/index/search/NestedHelperTests.java b/server/src/test/java/org/opensearch/index/search/NestedHelperTests.java index f7f921e824490..173cbf6b540a6 100644 --- a/server/src/test/java/org/opensearch/index/search/NestedHelperTests.java +++ b/server/src/test/java/org/opensearch/index/search/NestedHelperTests.java @@ -42,6 +42,7 @@ import org.apache.lucene.search.Query; import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.join.ScoreMode; +import org.opensearch.common.lucene.search.Queries; import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.xcontent.XContentBuilder; @@ -52,6 +53,8 @@ import org.opensearch.index.query.NestedQueryBuilder; import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.OpenSearchSingleNodeTestCase; import java.io.IOException; @@ -325,7 +328,10 @@ public void testNested() throws IOException { NestedQueryBuilder queryBuilder = new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.Avg); OpenSearchToParentBlockJoinQuery query = (OpenSearchToParentBlockJoinQuery) queryBuilder.toQuery(context); - Query expectedChildQuery = new BooleanQuery.Builder().add(new MatchAllDocsQuery(), Occur.MUST) + Query expectedChildQuery = new BooleanQuery.Builder().add( + new ApproximateScoreQuery(Queries.newMatchAllQuery(), new ApproximateMatchAllQuery()), + Occur.MUST + ) // we automatically add a filter since the inner query might match non-nested docs .add(new TermQuery(new Term(NestedPathFieldMapper.NAME, "nested1")), Occur.FILTER) .build(); diff --git a/server/src/test/java/org/opensearch/search/sort/GeoDistanceSortBuilderTests.java b/server/src/test/java/org/opensearch/search/sort/GeoDistanceSortBuilderTests.java index 385ced3655116..bca1ba13d11c3 100644 --- a/server/src/test/java/org/opensearch/search/sort/GeoDistanceSortBuilderTests.java +++ b/server/src/test/java/org/opensearch/search/sort/GeoDistanceSortBuilderTests.java @@ -59,7 +59,10 @@ import org.opensearch.index.query.RangeQueryBuilder; import org.opensearch.search.DocValueFormat; import org.opensearch.search.MultiValueMode; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.test.geo.RandomGeoGenerator; +import org.hamcrest.CoreMatchers; import java.io.IOException; import java.util.ArrayList; @@ -544,7 +547,10 @@ public void testBuildNested() throws IOException { XFieldComparatorSource comparatorSource = (XFieldComparatorSource) sortField.getComparatorSource(); Nested nested = comparatorSource.nested(); assertNotNull(nested); - assertEquals(new MatchAllDocsQuery(), nested.getInnerQuery()); + assertThat(nested.getInnerQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) nested.getInnerQuery(); + assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class)); sortBuilder = new GeoDistanceSortBuilder("fieldName", 1.0, 1.0).setNestedPath("path"); sortField = sortBuilder.build(shardContextMock).field; @@ -561,7 +567,10 @@ public void testBuildNested() throws IOException { comparatorSource = (XFieldComparatorSource) sortField.getComparatorSource(); nested = comparatorSource.nested(); assertNotNull(nested); - assertEquals(new MatchAllDocsQuery(), nested.getInnerQuery()); + assertThat(nested.getInnerQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class)); + approxQuery = (ApproximateScoreQuery) nested.getInnerQuery(); + assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class)); // if nested path is missing, we omit any filter and return a regular SortField // (LatLonSortField) diff --git a/server/src/test/java/org/opensearch/search/sort/ScriptSortBuilderTests.java b/server/src/test/java/org/opensearch/search/sort/ScriptSortBuilderTests.java index a124fdfeeb508..1660f8d73c323 100644 --- a/server/src/test/java/org/opensearch/search/sort/ScriptSortBuilderTests.java +++ b/server/src/test/java/org/opensearch/search/sort/ScriptSortBuilderTests.java @@ -54,7 +54,10 @@ import org.opensearch.script.ScriptType; import org.opensearch.search.DocValueFormat; import org.opensearch.search.MultiValueMode; +import org.opensearch.search.approximate.ApproximateMatchAllQuery; +import org.opensearch.search.approximate.ApproximateScoreQuery; import org.opensearch.search.sort.ScriptSortBuilder.ScriptSortType; +import org.hamcrest.CoreMatchers; import java.io.IOException; import java.util.Collections; @@ -336,7 +339,10 @@ public void testBuildNested() throws IOException { XFieldComparatorSource comparatorSource = (XFieldComparatorSource) sortField.getComparatorSource(); Nested nested = comparatorSource.nested(); assertNotNull(nested); - assertEquals(new MatchAllDocsQuery(), nested.getInnerQuery()); + assertThat(nested.getInnerQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class)); + ApproximateScoreQuery approxQuery = (ApproximateScoreQuery) nested.getInnerQuery(); + assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class)); sortBuilder = new ScriptSortBuilder(mockScript(MOCK_SCRIPT_NAME), ScriptSortType.NUMBER).setNestedPath("path"); sortField = sortBuilder.build(shardContextMock).field; @@ -353,7 +359,10 @@ public void testBuildNested() throws IOException { comparatorSource = (XFieldComparatorSource) sortField.getComparatorSource(); nested = comparatorSource.nested(); assertNotNull(nested); - assertEquals(new MatchAllDocsQuery(), nested.getInnerQuery()); + assertThat(nested.getInnerQuery(), CoreMatchers.instanceOf(ApproximateScoreQuery.class)); + approxQuery = (ApproximateScoreQuery) nested.getInnerQuery(); + assertThat(approxQuery.getOriginalQuery(), CoreMatchers.instanceOf(MatchAllDocsQuery.class)); + assertThat(approxQuery.getApproximationQuery(), CoreMatchers.instanceOf(ApproximateMatchAllQuery.class)); // if nested path is missing, we omit nested element in the comparator sortBuilder = new ScriptSortBuilder(mockScript(MOCK_SCRIPT_NAME), ScriptSortType.NUMBER).setNestedFilter( From c677dc7c71a7844323397675af554641bb0e5121 Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Thu, 10 Apr 2025 16:42:22 -0700 Subject: [PATCH 5/7] Fix more tests Signed-off-by: Harsha Vamsi Kalluri --- CHANGELOG.md | 1 + .../java/org/opensearch/search/sort/FieldSortIT.java | 3 +-- .../org/opensearch/validate/SimpleValidateQueryIT.java | 5 ++++- .../search/approximate/ApproximateMatchAllQuery.java | 8 ++++++-- .../search/approximate/ApproximateScoreQuery.java | 3 --- .../java/org/opensearch/search/sort/FieldSortBuilder.java | 1 + 6 files changed, 13 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bc7b66d47eb2..1dd96d71c5c6e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add SearchService and Search GRPC endpoint ([#17830](https://github.com/opensearch-project/OpenSearch/pull/17830)) - Add update and delete support in pull-based ingestion ([#17822](https://github.com/opensearch-project/OpenSearch/pull/17822)) - Allow maxPollSize and pollTimeout in IngestionSource to be configurable ([#17863](https://github.com/opensearch-project/OpenSearch/pull/17863)) +- Add `ApproximateMatchAllQuery` that targets match_all queries and approximates sorts ([#17772](https://github.com/opensearch-project/OpenSearch/pull/17772)) ### Changed - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912)) diff --git a/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java b/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java index 2017557852865..53692292e4169 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java @@ -1145,11 +1145,10 @@ public void testSortMissingNumbersMinMax() throws Exception { .get(); assertNoFailures(searchResponse); - assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(3L)); + assertThat(searchResponse.getHits().getTotalHits().value(), equalTo(2L)); assertThat(searchResponse.getHits().getAt(0).getId(), equalTo("1")); // The order here could be unstable (depends on document order) since missing == field value assertThat(searchResponse.getHits().getAt(1).getId(), is(oneOf("3", "2"))); - assertThat(searchResponse.getHits().getAt(2).getId(), is(oneOf("2", "3"))); logger.info("--> sort with missing _last"); searchResponse = client().prepareSearch() diff --git a/server/src/internalClusterTest/java/org/opensearch/validate/SimpleValidateQueryIT.java b/server/src/internalClusterTest/java/org/opensearch/validate/SimpleValidateQueryIT.java index 22647c59d42cf..735d201a94f2a 100644 --- a/server/src/internalClusterTest/java/org/opensearch/validate/SimpleValidateQueryIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/validate/SimpleValidateQueryIT.java @@ -294,7 +294,10 @@ public void testExplainNoQuery() { assertThat(validateQueryResponse.isValid(), equalTo(true)); assertThat(validateQueryResponse.getQueryExplanation().size(), equalTo(1)); assertThat(validateQueryResponse.getQueryExplanation().get(0).getIndex(), equalTo("test")); - assertThat(validateQueryResponse.getQueryExplanation().get(0).getExplanation(), equalTo("*:*")); + assertThat( + validateQueryResponse.getQueryExplanation().get(0).getExplanation(), + equalTo("ApproximateScoreQuery(originalQuery=*:*, approximationQuery=Approximate(*:*))") + ); } public void testExplainFilteredAlias() { diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java index 6a96b512acc48..d27c2af1f56bf 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -37,11 +37,15 @@ protected boolean canApproximate(SearchContext context) { return false; } - if (context.request() != null && context.request().source() != null) { + if (context.request() != null + && context.request().source() != null + && context.innerHits() == null + && context.aggregations() == null) { FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); if (primarySortField != null && primarySortField.missing() == null - && !primarySortField.fieldName().equals(FieldSortBuilder.DOC_FIELD_NAME)) { + && !primarySortField.fieldName().equals(FieldSortBuilder.DOC_FIELD_NAME) + && !primarySortField.fieldName().equals(FieldSortBuilder.ID_FIELD_NAME)) { MappedFieldType mappedFieldType = context.getQueryShardContext().fieldMapper(primarySortField.fieldName()); if (mappedFieldType == null) { return false; 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 6b39606620716..cf9ead3f088c2 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateScoreQuery.java @@ -51,9 +51,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { } public void setContext(SearchContext context) { - if (resolvedQuery != null) { - throw new IllegalStateException("Query already resolved, duplicate call to setContext"); - } resolvedQuery = approximationQuery.canApproximate(context) ? approximationQuery : originalQuery; }; diff --git a/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java b/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java index 9825b2cbbe08e..29de7fa66b78e 100644 --- a/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java +++ b/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java @@ -100,6 +100,7 @@ public class FieldSortBuilder extends SortBuilder implements W * special field name to sort by index order */ public static final String DOC_FIELD_NAME = "_doc"; + public static final String ID_FIELD_NAME = "_id"; private static final SortFieldAndFormat SORT_DOC = new SortFieldAndFormat(new SortField(null, SortField.Type.DOC), DocValueFormat.RAW); private static final SortFieldAndFormat SORT_DOC_REVERSE = new SortFieldAndFormat( new SortField(null, SortField.Type.DOC, true), From 05168da00d1df4b389fd23d4e23f07fc0af44d54 Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Fri, 11 Apr 2025 09:32:01 -0700 Subject: [PATCH 6/7] Fix more tests Signed-off-by: Harsha Vamsi Kalluri --- .../search/approximate/ApproximateMatchAllQuery.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java index d27c2af1f56bf..55d011d578b66 100644 --- a/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java +++ b/server/src/main/java/org/opensearch/search/approximate/ApproximateMatchAllQuery.java @@ -37,10 +37,7 @@ protected boolean canApproximate(SearchContext context) { return false; } - if (context.request() != null - && context.request().source() != null - && context.innerHits() == null - && context.aggregations() == null) { + if (context.request() != null && context.request().source() != null && context.innerHits().getInnerHits().isEmpty()) { FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source()); if (primarySortField != null && primarySortField.missing() == null From 93877eff8571dd5cba7c9ab431cfaa0dcefb851e Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Fri, 11 Apr 2025 10:59:25 -0700 Subject: [PATCH 7/7] Fix backward tests Signed-off-by: Harsha Vamsi Kalluri --- .../rest-api-spec/test/indices.validate_query/10_basic.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml index a114ed48741cf..a7f6c2adea921 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.validate_query/10_basic.yml @@ -13,8 +13,8 @@ setup: --- "Validate query api": - skip: - version: ' - 7.6.99' - reason: message changed in 7.7.0 + version: ' - 2.99.99' + reason: message changed in 3.0.0 - do: indices.validate_query: