From 579f2aaef81cdee80028fceba1f509dc69e108c6 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Fri, 30 Aug 2024 14:07:43 -0700 Subject: [PATCH 1/2] Add support for tracking failures at query group level (#15527) * add workload managementRequestFailureListener Signed-off-by: Kaushal Kumar * add unit tests Signed-off-by: Kaushal Kumar * add CHANGELOG Signed-off-by: Kaushal Kumar * add missing javadoc Signed-off-by: Kaushal Kumar * refactor Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * rename listener instance Signed-off-by: Kaushal Kumar --------- Signed-off-by: Kaushal Kumar --- CHANGELOG.md | 1 + .../main/java/org/opensearch/node/Node.java | 15 +- .../org/opensearch/wlm/QueryGroupService.java | 50 ++++- .../java/org/opensearch/wlm/ResourceType.java | 7 + ...> QueryGroupRequestOperationListener.java} | 13 +- .../opensearch/wlm/stats/QueryGroupState.java | 11 +- .../opensearch/wlm/stats/QueryGroupStats.java | 65 +++++- ...eryGroupRequestOperationListenerTests.java | 187 ++++++++++++++++++ ...equestRejectionOperationListenerTests.java | 53 ----- 9 files changed, 328 insertions(+), 74 deletions(-) rename server/src/main/java/org/opensearch/wlm/listeners/{QueryGroupRequestRejectionOperationListener.java => QueryGroupRequestOperationListener.java} (64%) create mode 100644 server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java delete mode 100644 server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListenerTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 9381aeee9865f..701eb72da62e4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Adding translog durability validation in index templates ([#15494](https://github.com/opensearch-project/OpenSearch/pull/15494)) - Add index creation using the context field ([#15290](https://github.com/opensearch-project/OpenSearch/pull/15290)) - [Reader Writer Separation] Add searchOnly replica routing configuration ([#15410](https://github.com/opensearch-project/OpenSearch/pull/15410)) +- [Workload Management] Add query group level failure tracking ([#15227](https://github.com/opensearch-project/OpenSearch/pull/15527)) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index ea656af6110e5..6373621c1143f 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -269,7 +269,7 @@ import org.opensearch.watcher.ResourceWatcherService; import org.opensearch.wlm.QueryGroupService; import org.opensearch.wlm.WorkloadManagementTransportInterceptor; -import org.opensearch.wlm.listeners.QueryGroupRequestRejectionOperationListener; +import org.opensearch.wlm.listeners.QueryGroupRequestOperationListener; import javax.net.ssl.SNIHostName; @@ -1019,11 +1019,12 @@ protected Node( List identityAwarePlugins = pluginsService.filterPlugins(IdentityAwarePlugin.class); identityService.initializeIdentityAwarePlugins(identityAwarePlugins); - final QueryGroupRequestRejectionOperationListener queryGroupRequestRejectionListener = - new QueryGroupRequestRejectionOperationListener( - new QueryGroupService(), // We will need to replace this with actual instance of the queryGroupService - threadPool - ); + final QueryGroupService queryGroupService = new QueryGroupService(); // We will need to replace this with actual instance of the + // queryGroupService + final QueryGroupRequestOperationListener queryGroupRequestOperationListener = new QueryGroupRequestOperationListener( + queryGroupService, + threadPool + ); // register all standard SearchRequestOperationsCompositeListenerFactory to the SearchRequestOperationsCompositeListenerFactory final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory = @@ -1033,7 +1034,7 @@ protected Node( searchRequestStats, searchRequestSlowLog, searchTaskRequestOperationsListener, - queryGroupRequestRejectionListener + queryGroupRequestOperationListener ), pluginComponents.stream() .filter(p -> p instanceof SearchRequestOperationsListener) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index 97c4e5169b4ed..6545598dd9951 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -9,11 +9,59 @@ package org.opensearch.wlm; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.wlm.stats.QueryGroupState; +import org.opensearch.wlm.stats.QueryGroupStats; +import org.opensearch.wlm.stats.QueryGroupStats.QueryGroupStatsHolder; + +import java.util.HashMap; +import java.util.Map; /** - * This is stub at this point in time and will be replace by an acutal one in couple of days + * As of now this is a stub and main implementation PR will be raised soon.Coming PR will collate these changes with core QueryGroupService changes */ public class QueryGroupService { + // This map does not need to be concurrent since we will process the cluster state change serially and update + // this map with new additions and deletions of entries. QueryGroupState is thread safe + private final Map queryGroupStateMap; + + public QueryGroupService() { + this(new HashMap<>()); + } + + public QueryGroupService(Map queryGroupStateMap) { + this.queryGroupStateMap = queryGroupStateMap; + } + + /** + * updates the failure stats for the query group + * @param queryGroupId query group identifier + */ + public void incrementFailuresFor(final String queryGroupId) { + QueryGroupState queryGroupState = queryGroupStateMap.get(queryGroupId); + // This can happen if the request failed for a deleted query group + // or new queryGroup is being created and has not been acknowledged yet + if (queryGroupState == null) { + return; + } + queryGroupState.failures.inc(); + } + + /** + * + * @return node level query group stats + */ + public QueryGroupStats nodeStats() { + final Map statsHolderMap = new HashMap<>(); + for (Map.Entry queryGroupsState : queryGroupStateMap.entrySet()) { + final String queryGroupId = queryGroupsState.getKey(); + final QueryGroupState currentState = queryGroupsState.getValue(); + + statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState)); + } + + return new QueryGroupStats(statsHolderMap); + } + /** * * @param queryGroupId query group identifier diff --git a/server/src/main/java/org/opensearch/wlm/ResourceType.java b/server/src/main/java/org/opensearch/wlm/ResourceType.java index c3f48f5f793ce..2e8da4f57f36c 100644 --- a/server/src/main/java/org/opensearch/wlm/ResourceType.java +++ b/server/src/main/java/org/opensearch/wlm/ResourceType.java @@ -14,6 +14,7 @@ import org.opensearch.tasks.Task; import java.io.IOException; +import java.util.List; import java.util.function.Function; /** @@ -30,6 +31,8 @@ public enum ResourceType { private final Function getResourceUsage; private final boolean statsEnabled; + private static List sortedValues = List.of(CPU, MEMORY); + ResourceType(String name, Function getResourceUsage, boolean statsEnabled) { this.name = name; this.getResourceUsage = getResourceUsage; @@ -71,4 +74,8 @@ public long getResourceUsage(Task task) { public boolean hasStatsEnabled() { return statsEnabled; } + + public static List getSortedValues() { + return sortedValues; + } } diff --git a/server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListener.java b/server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListener.java similarity index 64% rename from server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListener.java rename to server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListener.java index 89f6fe709667f..a2ce2b57bfe0f 100644 --- a/server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListener.java +++ b/server/src/main/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListener.java @@ -8,6 +8,7 @@ package org.opensearch.wlm.listeners; +import org.opensearch.action.search.SearchPhaseContext; import org.opensearch.action.search.SearchRequestContext; import org.opensearch.action.search.SearchRequestOperationsListener; import org.opensearch.threadpool.ThreadPool; @@ -15,14 +16,14 @@ import org.opensearch.wlm.QueryGroupTask; /** - * This listener is used to perform the rejections for incoming requests into a queryGroup + * This listener is used to listen for request lifecycle events for a queryGroup */ -public class QueryGroupRequestRejectionOperationListener extends SearchRequestOperationsListener { +public class QueryGroupRequestOperationListener extends SearchRequestOperationsListener { private final QueryGroupService queryGroupService; private final ThreadPool threadPool; - public QueryGroupRequestRejectionOperationListener(QueryGroupService queryGroupService, ThreadPool threadPool) { + public QueryGroupRequestOperationListener(QueryGroupService queryGroupService, ThreadPool threadPool) { this.queryGroupService = queryGroupService; this.threadPool = threadPool; } @@ -36,4 +37,10 @@ protected void onRequestStart(SearchRequestContext searchRequestContext) { final String queryGroupId = threadPool.getThreadContext().getHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER); queryGroupService.rejectIfNeeded(queryGroupId); } + + @Override + protected void onRequestFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) { + final String queryGroupId = threadPool.getThreadContext().getHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER); + queryGroupService.incrementFailuresFor(queryGroupId); + } } diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java index 93cfcea697c43..376d34dd7c8ca 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java @@ -31,7 +31,7 @@ public class QueryGroupState { /** * this will track the cumulative failures in a query group */ - final CounterMetric failures = new CounterMetric(); + public final CounterMetric failures = new CounterMetric(); /** * This will track total number of cancellations in the query group due to all resource type breaches @@ -95,9 +95,18 @@ public static class ResourceTypeState { final ResourceType resourceType; final CounterMetric cancellations = new CounterMetric(); final CounterMetric rejections = new CounterMetric(); + private double lastRecordedUsage = 0; public ResourceTypeState(ResourceType resourceType) { this.resourceType = resourceType; } + + public void setLastRecordedUsage(double recordedUsage) { + lastRecordedUsage = recordedUsage; + } + + public double getLastRecordedUsage() { + return lastRecordedUsage; + } } } diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java index d39bf104332da..9fc7039cd1852 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -14,8 +14,12 @@ import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.stats.QueryGroupState.ResourceTypeState; import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Objects; @@ -52,7 +56,11 @@ public void writeTo(StreamOutput out) throws IOException { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject("query_groups"); - for (Map.Entry queryGroupStats : stats.entrySet()) { + // to keep the toXContent consistent + List> entryList = new ArrayList<>(stats.entrySet()); + entryList.sort((k1, k2) -> k1.getKey().compareTo(k2.getKey())); + + for (Map.Entry queryGroupStats : entryList) { builder.startObject(queryGroupStats.getKey()); queryGroupStats.getValue().toXContent(builder, params); builder.endObject(); @@ -83,11 +91,14 @@ public static class QueryGroupStatsHolder implements ToXContentObject, Writeable public static final String REJECTIONS = "rejections"; public static final String TOTAL_CANCELLATIONS = "total_cancellations"; public static final String FAILURES = "failures"; - private final long completions; - private final long rejections; - private final long failures; - private final long totalCancellations; - private final Map resourceStats; + private long completions; + private long rejections; + private long failures; + private long totalCancellations; + private Map resourceStats; + + // this is needed to support the factory method + public QueryGroupStatsHolder() {} public QueryGroupStatsHolder( long completions, @@ -111,6 +122,28 @@ public QueryGroupStatsHolder(StreamInput in) throws IOException { this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } + /** + * static factory method to convert {@link QueryGroupState} into {@link QueryGroupStatsHolder} + * @param queryGroupState which needs to be converted + * @return QueryGroupStatsHolder object + */ + public static QueryGroupStatsHolder from(QueryGroupState queryGroupState) { + final QueryGroupStatsHolder statsHolder = new QueryGroupStatsHolder(); + + Map resourceStatsMap = new HashMap<>(); + + for (Map.Entry resourceTypeStateEntry : queryGroupState.getResourceState().entrySet()) { + resourceStatsMap.put(resourceTypeStateEntry.getKey(), ResourceStats.from(resourceTypeStateEntry.getValue())); + } + + statsHolder.completions = queryGroupState.getCompletions(); + statsHolder.rejections = queryGroupState.getTotalRejections(); + statsHolder.failures = queryGroupState.getFailures(); + statsHolder.totalCancellations = queryGroupState.getTotalCancellations(); + statsHolder.resourceStats = resourceStatsMap; + return statsHolder; + } + /** * Writes the {@param statsHolder} to {@param out} * @param out StreamOutput @@ -136,9 +169,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field(REJECTIONS, rejections); builder.field(FAILURES, failures); builder.field(TOTAL_CANCELLATIONS, totalCancellations); - for (Map.Entry resourceStat : resourceStats.entrySet()) { - ResourceType resourceType = resourceStat.getKey(); - ResourceStats resourceStats1 = resourceStat.getValue(); + + for (ResourceType resourceType : ResourceType.getSortedValues()) { + ResourceStats resourceStats1 = resourceStats.get(resourceType); + if (resourceStats1 == null) continue; builder.startObject(resourceType.getName()); resourceStats1.toXContent(builder, params); builder.endObject(); @@ -187,6 +221,19 @@ public ResourceStats(StreamInput in) throws IOException { this.rejections = in.readVLong(); } + /** + * static factory method to convert {@link ResourceTypeState} into {@link ResourceStats} + * @param resourceTypeState which needs to be converted + * @return QueryGroupStatsHolder object + */ + public static ResourceStats from(ResourceTypeState resourceTypeState) { + return new ResourceStats( + resourceTypeState.getLastRecordedUsage(), + resourceTypeState.cancellations.count(), + resourceTypeState.rejections.count() + ); + } + /** * Writes the {@param stats} to {@param out} * @param out StreamOutput diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java new file mode 100644 index 0000000000000..0307ff623c408 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java @@ -0,0 +1,187 @@ +/* + * 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.wlm.listeners; + +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupService; +import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.stats.QueryGroupState; +import org.opensearch.wlm.stats.QueryGroupStats; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; + +public class QueryGroupRequestOperationListenerTests extends OpenSearchTestCase { + public static final int ITERATIONS = 20; + ThreadPool testThreadPool; + QueryGroupService queryGroupService; + + Map queryGroupStateMap; + String testQueryGroupId; + QueryGroupRequestOperationListener sut; + + public void setUp() throws Exception { + super.setUp(); + queryGroupStateMap = new HashMap<>(); + testQueryGroupId = "safjgagnakg-3r3fads"; + testThreadPool = new TestThreadPool("RejectionTestThreadPool"); + queryGroupService = mock(QueryGroupService.class); + sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); + } + + public void tearDown() throws Exception { + super.tearDown(); + testThreadPool.shutdown(); + } + + public void testRejectionCase() { + final String testQueryGroupId = "asdgasgkajgkw3141_3rt4t"; + testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, testQueryGroupId); + doThrow(OpenSearchRejectedExecutionException.class).when(queryGroupService).rejectIfNeeded(testQueryGroupId); + assertThrows(OpenSearchRejectedExecutionException.class, () -> sut.onRequestStart(null)); + } + + public void testNonRejectionCase() { + final String testQueryGroupId = "asdgasgkajgkw3141_3rt4t"; + testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, testQueryGroupId); + doNothing().when(queryGroupService).rejectIfNeeded(testQueryGroupId); + + sut.onRequestStart(null); + } + + public void testValidQueryGroupRequestFailure() throws IOException { + + QueryGroupStats expectedStats = new QueryGroupStats( + Map.of( + testQueryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 1, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ) + ) + ); + + assertSuccess(testQueryGroupId, queryGroupStateMap, expectedStats, testQueryGroupId); + } + + public void testMultiThreadedValidQueryGroupRequestFailures() { + + queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); + + queryGroupService = new QueryGroupService(queryGroupStateMap); + + sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); + + List threads = new ArrayList<>(); + for (int i = 0; i < ITERATIONS; i++) { + threads.add(new Thread(() -> { + try (ThreadContext.StoredContext currentContext = testThreadPool.getThreadContext().stashContext()) { + testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, testQueryGroupId); + sut.onRequestFailure(null, null); + } + })); + } + + threads.forEach(Thread::start); + threads.forEach(th -> { + try { + th.join(); + } catch (InterruptedException ignored) { + + } + }); + + QueryGroupStats actualStats = queryGroupService.nodeStats(); + + QueryGroupStats expectedStats = new QueryGroupStats( + Map.of( + testQueryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + ITERATIONS, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ) + ) + ); + + assertEquals(expectedStats, actualStats); + } + + public void testInvalidQueryGroupFailure() throws IOException { + QueryGroupStats expectedStats = new QueryGroupStats( + Map.of( + testQueryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ) + ) + ); + + assertSuccess(testQueryGroupId, queryGroupStateMap, expectedStats, "dummy-invalid-qg-id"); + + } + + private void assertSuccess( + String testQueryGroupId, + Map queryGroupStateMap, + QueryGroupStats expectedStats, + String threadContextQG_Id + ) { + + try (ThreadContext.StoredContext currentContext = testThreadPool.getThreadContext().stashContext()) { + testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, threadContextQG_Id); + queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); + + queryGroupService = new QueryGroupService(queryGroupStateMap); + + sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); + sut.onRequestFailure(null, null); + + QueryGroupStats actualStats = queryGroupService.nodeStats(); + assertEquals(expectedStats, actualStats); + } + + } +} diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListenerTests.java deleted file mode 100644 index 19e82aca26153..0000000000000 --- a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestRejectionOperationListenerTests.java +++ /dev/null @@ -1,53 +0,0 @@ -/* - * 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.wlm.listeners; - -import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; -import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.threadpool.TestThreadPool; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.wlm.QueryGroupService; -import org.opensearch.wlm.QueryGroupTask; - -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; - -public class QueryGroupRequestRejectionOperationListenerTests extends OpenSearchTestCase { - ThreadPool testThreadPool; - QueryGroupService queryGroupService; - QueryGroupRequestRejectionOperationListener sut; - - public void setUp() throws Exception { - super.setUp(); - testThreadPool = new TestThreadPool("RejectionTestThreadPool"); - queryGroupService = mock(QueryGroupService.class); - sut = new QueryGroupRequestRejectionOperationListener(queryGroupService, testThreadPool); - } - - public void tearDown() throws Exception { - super.tearDown(); - testThreadPool.shutdown(); - } - - public void testRejectionCase() { - final String testQueryGroupId = "asdgasgkajgkw3141_3rt4t"; - testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, testQueryGroupId); - doThrow(OpenSearchRejectedExecutionException.class).when(queryGroupService).rejectIfNeeded(testQueryGroupId); - assertThrows(OpenSearchRejectedExecutionException.class, () -> sut.onRequestStart(null)); - } - - public void testNonRejectionCase() { - final String testQueryGroupId = "asdgasgkajgkw3141_3rt4t"; - testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, testQueryGroupId); - doNothing().when(queryGroupService).rejectIfNeeded(testQueryGroupId); - - sut.onRequestStart(null); - } -} From 839ba0bca05f8a9aefdbabd023dc75741265bd6c Mon Sep 17 00:00:00 2001 From: David Zane <38449481+dzane17@users.noreply.github.com> Date: Fri, 30 Aug 2024 14:58:59 -0700 Subject: [PATCH 2/2] Add fieldType to AbstractQueryBuilder and SortBuilder (#15328) * Add fieldType to AbstractQueryBuilder and FieldSortBuilder Signed-off-by: David Zane * Add fieldType to AbstractQueryBuilder and SortBuilder Signed-off-by: David Zane --------- Signed-off-by: David Zane Signed-off-by: Ankit Jain Co-authored-by: Ankit Jain --- CHANGELOG.md | 1 + .../index/query/RankFeatureQueryBuilder.java | 5 ++ .../join/query/HasChildQueryBuilder.java | 5 ++ .../join/query/HasParentQueryBuilder.java | 5 ++ .../join/query/ParentIdQueryBuilder.java | 5 ++ .../percolator/PercolateQueryBuilder.java | 5 ++ .../index/mapper/FieldTypeLookup.java | 3 +- .../index/mapper/MapperService.java | 3 +- .../index/query/AbstractQueryBuilder.java | 23 ++++++ .../index/query/BoolQueryBuilder.java | 5 ++ .../index/query/BoostingQueryBuilder.java | 5 ++ .../query/ConstantScoreQueryBuilder.java | 5 ++ .../index/query/DisMaxQueryBuilder.java | 5 ++ .../query/DistanceFeatureQueryBuilder.java | 3 +- .../index/query/IdsQueryBuilder.java | 5 ++ .../index/query/IntervalQueryBuilder.java | 5 ++ .../index/query/MatchAllQueryBuilder.java | 5 ++ .../index/query/MatchNoneQueryBuilder.java | 5 ++ .../index/query/MoreLikeThisQueryBuilder.java | 5 ++ .../index/query/MultiMatchQueryBuilder.java | 5 ++ .../index/query/NestedQueryBuilder.java | 5 ++ .../index/query/QueryShardContext.java | 16 ++++ .../index/query/QueryStringQueryBuilder.java | 5 ++ .../index/query/ScriptQueryBuilder.java | 5 ++ .../index/query/SimpleQueryStringBuilder.java | 5 ++ .../query/SpanContainingQueryBuilder.java | 5 ++ .../index/query/SpanFirstQueryBuilder.java | 5 ++ .../query/SpanMultiTermQueryBuilder.java | 5 ++ .../index/query/SpanNearQueryBuilder.java | 5 ++ .../index/query/SpanNotQueryBuilder.java | 5 ++ .../index/query/SpanOrQueryBuilder.java | 5 ++ .../index/query/SpanWithinQueryBuilder.java | 5 ++ .../index/query/TermsSetQueryBuilder.java | 5 ++ .../index/query/WrapperQueryBuilder.java | 5 ++ .../FunctionScoreQueryBuilder.java | 5 ++ .../ScriptScoreQueryBuilder.java | 5 ++ .../search/sort/FieldSortBuilder.java | 5 ++ .../search/sort/ScoreSortBuilder.java | 5 ++ .../search/sort/ScriptSortBuilder.java | 5 ++ .../opensearch/search/sort/SortBuilder.java | 30 +++++++- .../index/query/QueryShardContextTests.java | 75 +++++++++++++++++++ .../index/query/plugin/DummyQueryBuilder.java | 5 ++ .../opensearch/search/SearchServiceTests.java | 5 ++ .../search/sort/AbstractSortTestCase.java | 13 +++- .../search/sort/FieldSortBuilderTests.java | 38 ++++++++++ .../search/sort/plugin/CustomSortBuilder.java | 5 ++ 46 files changed, 379 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 701eb72da62e4..5e780538bf147 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Workload Management] Add rejection logic for co-ordinator and shard level requests ([#15428](https://github.com/opensearch-project/OpenSearch/pull/15428))) - Adding translog durability validation in index templates ([#15494](https://github.com/opensearch-project/OpenSearch/pull/15494)) - Add index creation using the context field ([#15290](https://github.com/opensearch-project/OpenSearch/pull/15290)) +- Add fieldType to AbstractQueryBuilder and FieldSortBuilder ([#15328](https://github.com/opensearch-project/OpenSearch/pull/15328))) - [Reader Writer Separation] Add searchOnly replica routing configuration ([#15410](https://github.com/opensearch-project/OpenSearch/pull/15410)) - [Workload Management] Add query group level failure tracking ([#15227](https://github.com/opensearch-project/OpenSearch/pull/15527)) diff --git a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java index 13591d0782ea2..53bafa1346711 100644 --- a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java +++ b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java @@ -400,6 +400,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + @Override protected Query doToQuery(QueryShardContext context) throws IOException { final MappedFieldType ft = context.fieldMapper(field); diff --git a/modules/parent-join/src/main/java/org/opensearch/join/query/HasChildQueryBuilder.java b/modules/parent-join/src/main/java/org/opensearch/join/query/HasChildQueryBuilder.java index e930780613ed6..837cf8981fcee 100644 --- a/modules/parent-join/src/main/java/org/opensearch/join/query/HasChildQueryBuilder.java +++ b/modules/parent-join/src/main/java/org/opensearch/join/query/HasChildQueryBuilder.java @@ -264,6 +264,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static HasChildQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String childType = null; diff --git a/modules/parent-join/src/main/java/org/opensearch/join/query/HasParentQueryBuilder.java b/modules/parent-join/src/main/java/org/opensearch/join/query/HasParentQueryBuilder.java index d296a7b0141ff..e1438f0713034 100644 --- a/modules/parent-join/src/main/java/org/opensearch/join/query/HasParentQueryBuilder.java +++ b/modules/parent-join/src/main/java/org/opensearch/join/query/HasParentQueryBuilder.java @@ -233,6 +233,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static HasParentQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String parentType = null; diff --git a/modules/parent-join/src/main/java/org/opensearch/join/query/ParentIdQueryBuilder.java b/modules/parent-join/src/main/java/org/opensearch/join/query/ParentIdQueryBuilder.java index bfc01bf151a9c..605f2abcabdc4 100644 --- a/modules/parent-join/src/main/java/org/opensearch/join/query/ParentIdQueryBuilder.java +++ b/modules/parent-join/src/main/java/org/opensearch/join/query/ParentIdQueryBuilder.java @@ -130,6 +130,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static ParentIdQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String type = null; diff --git a/modules/percolator/src/main/java/org/opensearch/percolator/PercolateQueryBuilder.java b/modules/percolator/src/main/java/org/opensearch/percolator/PercolateQueryBuilder.java index 6933bfbef4666..38e941fc42b87 100644 --- a/modules/percolator/src/main/java/org/opensearch/percolator/PercolateQueryBuilder.java +++ b/modules/percolator/src/main/java/org/opensearch/percolator/PercolateQueryBuilder.java @@ -356,6 +356,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>(NAME, args -> { String field = (String) args[0]; BytesReference document = (BytesReference) args[1]; diff --git a/server/src/main/java/org/opensearch/index/mapper/FieldTypeLookup.java b/server/src/main/java/org/opensearch/index/mapper/FieldTypeLookup.java index 7a9a0c5fb5428..dfedbd23ae5ca 100644 --- a/server/src/main/java/org/opensearch/index/mapper/FieldTypeLookup.java +++ b/server/src/main/java/org/opensearch/index/mapper/FieldTypeLookup.java @@ -101,7 +101,8 @@ class FieldTypeLookup implements Iterable { } /** - * Returns the mapped field type for the given field name. + * Returns the {@link MappedFieldType} for the given field name + * or null if the field name is not found. */ public MappedFieldType get(String field) { String concreteField = aliasToConcreteName.getOrDefault(field, field); diff --git a/server/src/main/java/org/opensearch/index/mapper/MapperService.java b/server/src/main/java/org/opensearch/index/mapper/MapperService.java index 530a3092a5aa7..400232ff164db 100644 --- a/server/src/main/java/org/opensearch/index/mapper/MapperService.java +++ b/server/src/main/java/org/opensearch/index/mapper/MapperService.java @@ -622,7 +622,8 @@ public DocumentMapperForType documentMapperWithAutoCreate() { } /** - * Given the full name of a field, returns its {@link MappedFieldType}. + * Given the full name of a field, returns its {@link MappedFieldType} + * or null if the field is not found. */ public MappedFieldType fieldType(String fullName) { return this.mapper == null ? null : this.mapper.fieldTypes().get(fullName); diff --git a/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java index 66c6ee115c3f0..2e9b856b4fdc5 100644 --- a/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/AbstractQueryBuilder.java @@ -74,6 +74,7 @@ public abstract class AbstractQueryBuilder> public static final ParseField BOOST_FIELD = new ParseField("boost"); protected String queryName; + protected String fieldType; protected float boost = DEFAULT_BOOST; protected AbstractQueryBuilder() { @@ -112,6 +113,27 @@ protected void printBoostAndQueryName(XContentBuilder builder) throws IOExceptio } } + /** + * Returns field name as String. + * Abstract method to be implemented by all child classes. + */ + public abstract String fieldName(); + + /** + * Default method for child classes which do not have a custom {@link #fieldName()} implementation. + */ + protected static String getDefaultFieldName() { + return null; + }; + + /** + * Returns field type as String for QueryBuilder classes which have a defined fieldName. + * Else returns null. + */ + public final String getFieldType() { + return fieldType; + }; + @Override public final Query toQuery(QueryShardContext context) throws IOException { Query query = doToQuery(context); @@ -125,6 +147,7 @@ public final Query toQuery(QueryShardContext context) throws IOException { context.addNamedQuery(queryName, query); } } + fieldType = context.getFieldTypeString(fieldName()); return query; } diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java index c44a7ef6a397c..fe93a5aec6950 100644 --- a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java @@ -270,6 +270,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + private static void doXArrayContent(ParseField field, List clauses, XContentBuilder builder, Params params) throws IOException { if (clauses.isEmpty()) { diff --git a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java index 1b52ae2f03605..185dbf60b5a43 100644 --- a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java @@ -151,6 +151,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static BoostingQueryBuilder fromXContent(XContentParser parser) throws IOException { QueryBuilder positiveQuery = null; boolean positiveQueryFound = false; diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java index b2764d29da80a..db92536fe9201 100644 --- a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java @@ -101,6 +101,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static ConstantScoreQueryBuilder fromXContent(XContentParser parser) throws IOException { QueryBuilder query = null; boolean queryFound = false; diff --git a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java index bd8ec62f6c43e..fe90be91a8f89 100644 --- a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java @@ -137,6 +137,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static DisMaxQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; float tieBreaker = DisMaxQueryBuilder.DEFAULT_TIE_BREAKER; diff --git a/server/src/main/java/org/opensearch/index/query/DistanceFeatureQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/DistanceFeatureQueryBuilder.java index 1d9f0479c6b17..016b6357f4934 100644 --- a/server/src/main/java/org/opensearch/index/query/DistanceFeatureQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/DistanceFeatureQueryBuilder.java @@ -136,7 +136,8 @@ protected Query doToQuery(QueryShardContext context) throws IOException { return fieldType.distanceFeatureQuery(origin.origin(), pivot, 1.0f, context); } - String fieldName() { + @Override + public String fieldName() { return field; } diff --git a/server/src/main/java/org/opensearch/index/query/IdsQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/IdsQueryBuilder.java index d7ebdbff10adb..a47b37204fa55 100644 --- a/server/src/main/java/org/opensearch/index/query/IdsQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/IdsQueryBuilder.java @@ -132,6 +132,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + private static final ObjectParser PARSER = new ObjectParser<>(NAME, IdsQueryBuilder::new); static { diff --git a/server/src/main/java/org/opensearch/index/query/IntervalQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/IntervalQueryBuilder.java index 125035ea5e95a..3965bb47ad1b9 100644 --- a/server/src/main/java/org/opensearch/index/query/IntervalQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/IntervalQueryBuilder.java @@ -95,6 +95,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static IntervalQueryBuilder fromXContent(XContentParser parser) throws IOException { if (parser.nextToken() != XContentParser.Token.FIELD_NAME) { throw new ParsingException(parser.getTokenLocation(), "Expected [FIELD_NAME] but got [" + parser.currentToken() + "]"); 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..07e6737597977 100644 --- a/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MatchAllQueryBuilder.java @@ -72,6 +72,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + private static final ObjectParser PARSER = new ObjectParser<>(NAME, MatchAllQueryBuilder::new); static { diff --git a/server/src/main/java/org/opensearch/index/query/MatchNoneQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/MatchNoneQueryBuilder.java index 17e84bc785206..c07e5c9d77600 100644 --- a/server/src/main/java/org/opensearch/index/query/MatchNoneQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MatchNoneQueryBuilder.java @@ -71,6 +71,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static MatchNoneQueryBuilder fromXContent(XContentParser parser) throws IOException { String currentFieldName = null; XContentParser.Token token; diff --git a/server/src/main/java/org/opensearch/index/query/MoreLikeThisQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/MoreLikeThisQueryBuilder.java index e6472afef2215..141d4beb5f885 100644 --- a/server/src/main/java/org/opensearch/index/query/MoreLikeThisQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MoreLikeThisQueryBuilder.java @@ -792,6 +792,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static MoreLikeThisQueryBuilder fromXContent(XContentParser parser) throws IOException { // document inputs List fields = null; diff --git a/server/src/main/java/org/opensearch/index/query/MultiMatchQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/MultiMatchQueryBuilder.java index 6227e5d2fa806..8eb5beadf3e48 100644 --- a/server/src/main/java/org/opensearch/index/query/MultiMatchQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/MultiMatchQueryBuilder.java @@ -631,6 +631,11 @@ public void doXContent(XContentBuilder builder, Params params) throws IOExceptio builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static MultiMatchQueryBuilder fromXContent(XContentParser parser) throws IOException { Object value = null; Map fieldsBoosts = new HashMap<>(); diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java index 5908882472ce7..f5aab8c967b88 100644 --- a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java @@ -205,6 +205,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static NestedQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; ScoreMode scoreMode = ScoreMode.Avg; diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java index 91313092d8d28..b0ed049e6a1f4 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java +++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java @@ -364,6 +364,22 @@ public MappedFieldType fieldMapper(String name) { return failIfFieldMappingNotFound(name, mapperService.fieldType(name)); } + /** + * Returns field type as String for the given field name. + * If field is not mapped or mapperService is null, returns null. + */ + public String getFieldTypeString(String fieldName) { + if (fieldName != null) { + if (mapperService != null) { + MappedFieldType mappedFieldType = mapperService.fieldType(fieldName); + if (mappedFieldType != null) { + return mappedFieldType.typeName(); + } + } + } + return null; + } + public ObjectMapper getObjectMapper(String name) { return mapperService.getObjectMapper(name); } diff --git a/server/src/main/java/org/opensearch/index/query/QueryStringQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/QueryStringQueryBuilder.java index 3d8fbd5fc436d..e6148a23944d9 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryStringQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/QueryStringQueryBuilder.java @@ -647,6 +647,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static QueryStringQueryBuilder fromXContent(XContentParser parser) throws IOException { String currentFieldName = null; XContentParser.Token token; diff --git a/server/src/main/java/org/opensearch/index/query/ScriptQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/ScriptQueryBuilder.java index ded6fd0528c33..4d895bc8bb725 100644 --- a/server/src/main/java/org/opensearch/index/query/ScriptQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/ScriptQueryBuilder.java @@ -106,6 +106,11 @@ protected void doXContent(XContentBuilder builder, Params builderParams) throws builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static ScriptQueryBuilder fromXContent(XContentParser parser) throws IOException { // also, when caching, since its isCacheable is false, will result in loading all bit set... Script script = null; diff --git a/server/src/main/java/org/opensearch/index/query/SimpleQueryStringBuilder.java b/server/src/main/java/org/opensearch/index/query/SimpleQueryStringBuilder.java index 57ae7dd0ea5e9..5d5da6275601c 100644 --- a/server/src/main/java/org/opensearch/index/query/SimpleQueryStringBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SimpleQueryStringBuilder.java @@ -479,6 +479,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SimpleQueryStringBuilder fromXContent(XContentParser parser) throws IOException { String currentFieldName = null; String queryBody = null; diff --git a/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java index 32a19ea3e9b50..e7cc3c6f40855 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java @@ -117,6 +117,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanContainingQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String queryName = null; diff --git a/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java index bcbc64ddf386d..d8ad0c83ad0bd 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java @@ -120,6 +120,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanFirstQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; diff --git a/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java index 96d03c91964e3..f026dddfe3522 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java @@ -99,6 +99,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanMultiTermQueryBuilder fromXContent(XContentParser parser) throws IOException { String currentFieldName = null; MultiTermQueryBuilder subQuery = null; diff --git a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java index 30a1c29c29126..536024666dc38 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java @@ -167,6 +167,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanNearQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; int slop = DEFAULT_SLOP; diff --git a/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java index 59ec5b9d77fc8..75ccd028656c2 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java @@ -181,6 +181,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanNotQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; diff --git a/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java index fae1e318c66bd..2014d9b807b6f 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java @@ -115,6 +115,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanOrQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String queryName = null; diff --git a/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java index 4d5a6dde61a70..b1c0aebc1fb52 100644 --- a/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java @@ -122,6 +122,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static SpanWithinQueryBuilder fromXContent(XContentParser parser) throws IOException { float boost = AbstractQueryBuilder.DEFAULT_BOOST; String queryName = null; diff --git a/server/src/main/java/org/opensearch/index/query/TermsSetQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/TermsSetQueryBuilder.java index e2cf7384ecac7..d6b3ecd102c08 100644 --- a/server/src/main/java/org/opensearch/index/query/TermsSetQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/TermsSetQueryBuilder.java @@ -169,6 +169,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static TermsSetQueryBuilder fromXContent(XContentParser parser) throws IOException { XContentParser.Token token = parser.nextToken(); if (token != XContentParser.Token.FIELD_NAME) { diff --git a/server/src/main/java/org/opensearch/index/query/WrapperQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/WrapperQueryBuilder.java index 8a322b2f9e173..d08d50f17191b 100644 --- a/server/src/main/java/org/opensearch/index/query/WrapperQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/WrapperQueryBuilder.java @@ -131,6 +131,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static WrapperQueryBuilder fromXContent(XContentParser parser) throws IOException { XContentParser.Token token = parser.nextToken(); if (token != XContentParser.Token.FIELD_NAME) { diff --git a/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java index b3c797f11de6d..454bf62e837f9 100644 --- a/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java @@ -300,6 +300,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public FunctionScoreQueryBuilder setMinScore(float minScore) { this.minScore = minScore; return this; diff --git a/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java index fe9ad200d44f0..160616faed114 100644 --- a/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/functionscore/ScriptScoreQueryBuilder.java @@ -155,6 +155,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public ScriptScoreQueryBuilder setMinScore(float minScore) { this.minScore = minScore; return this; 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 5cecda1346b90..c93e37da6854d 100644 --- a/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java +++ b/server/src/main/java/org/opensearch/search/sort/FieldSortBuilder.java @@ -559,6 +559,11 @@ public BucketedSort buildBucketedSort(QueryShardContext context, int bucketSize, } } + @Override + public String fieldName() { + return fieldName; + } + private MappedFieldType resolveUnmappedType(QueryShardContext context) { if (unmappedType == null) { throw new QueryShardException(context, "No mapping found for [" + fieldName + "] in order to sort on"); diff --git a/server/src/main/java/org/opensearch/search/sort/ScoreSortBuilder.java b/server/src/main/java/org/opensearch/search/sort/ScoreSortBuilder.java index 1be49e5ca81ce..38774c0e10b71 100644 --- a/server/src/main/java/org/opensearch/search/sort/ScoreSortBuilder.java +++ b/server/src/main/java/org/opensearch/search/sort/ScoreSortBuilder.java @@ -160,6 +160,11 @@ protected float docValue() { }; } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + @Override public boolean equals(Object object) { if (this == object) { diff --git a/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java b/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java index bb1930eb3a953..0ffc2b960fcd3 100644 --- a/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java +++ b/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java @@ -335,6 +335,11 @@ public BucketedSort buildBucketedSort(QueryShardContext context, int bucketSize, return fieldComparatorSource(context).newBucketedSort(context.bigArrays(), order, DocValueFormat.RAW, bucketSize, extra); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + private IndexFieldData.XFieldComparatorSource fieldComparatorSource(QueryShardContext context) throws IOException { MultiValueMode valueMode = null; if (sortMode != null) { diff --git a/server/src/main/java/org/opensearch/search/sort/SortBuilder.java b/server/src/main/java/org/opensearch/search/sort/SortBuilder.java index a8c21e7311061..d6c5f2b5f3144 100644 --- a/server/src/main/java/org/opensearch/search/sort/SortBuilder.java +++ b/server/src/main/java/org/opensearch/search/sort/SortBuilder.java @@ -71,6 +71,7 @@ public abstract class SortBuilder> implements NamedWriteable, ToXContentObject, Rewriteable> { protected SortOrder order = SortOrder.ASC; + protected String fieldType; // parse fields common to more than one SortBuilder public static final ParseField ORDER_FIELD = new ParseField("order"); @@ -161,11 +162,17 @@ private static void parseCompoundSortField(XContentParser parser, List buildSort(List> sortBuilders, QueryShardContext context) throws IOException { List sortFields = new ArrayList<>(sortBuilders.size()); List sortFormats = new ArrayList<>(sortBuilders.size()); for (SortBuilder builder : sortBuilders) { - SortFieldAndFormat sf = builder.build(context); + SortFieldAndFormat sf = builder.buildHelper(context); sortFields.add(sf.field); sortFormats.add(sf.format); } @@ -287,4 +294,25 @@ protected static QueryBuilder parseNestedFilter(XContentParser parser) { public String toString() { return Strings.toString(MediaTypeRegistry.JSON, this, true, true); } + + /** + * Returns field name as String. + * Abstract method to be implemented by all child classes. + */ + public abstract String fieldName(); + + /** + * Default method for child classes which do not have a custom {@link #fieldName()} implementation. + */ + protected static String getDefaultFieldName() { + return null; + }; + + /** + * Returns field type as String for SortBuilder classes which have a defined fieldName. + * Else returns null. + */ + public final String getFieldType() { + return fieldType; + }; } diff --git a/server/src/test/java/org/opensearch/index/query/QueryShardContextTests.java b/server/src/test/java/org/opensearch/index/query/QueryShardContextTests.java index 12677edc8efa7..294e175ca670e 100644 --- a/server/src/test/java/org/opensearch/index/query/QueryShardContextTests.java +++ b/server/src/test/java/org/opensearch/index/query/QueryShardContextTests.java @@ -69,6 +69,8 @@ import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.TextFieldMapper; +import org.opensearch.index.mapper.TextSearchInfo; +import org.opensearch.index.mapper.ValueFetcher; import org.opensearch.search.lookup.LeafDocLookup; import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.search.lookup.SearchLookup; @@ -150,6 +152,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep } + @Override + public String fieldName() { + return getDefaultFieldName(); + } + @Override protected Query doToQuery(QueryShardContext context) throws IOException { throw new RuntimeException("boom"); @@ -479,4 +486,72 @@ public void collect(int doc) throws IOException { } } + public void testGetFieldTypeString() { + MapperService mapperService = mock(MapperService.class); + + Settings settings = Settings.builder() + .put("index.version.created", Version.CURRENT) + .put("index.number_of_shards", 1) + .put("index.number_of_replicas", 0) + .put(IndexMetadata.SETTING_INDEX_UUID, "uuid") + .build(); + IndexMetadata indexMetadata = new IndexMetadata.Builder("index").settings(settings).build(); + IndexSettings indexSettings = new IndexSettings(indexMetadata, settings); + QueryShardContext queryShardContext = new QueryShardContext( + 0, + indexSettings, + null, + null, + null, + mapperService, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ); + + MappedFieldType mappedFieldType = new MappedFieldType("field", true, false, true, TextSearchInfo.NONE, Collections.emptyMap()) { + @Override + public String typeName() { + return "long"; + } + + @Override + public ValueFetcher valueFetcher(QueryShardContext context, SearchLookup searchLookup, String format) { + throw new UnsupportedOperationException(); + } + + @Override + public Query termQuery(Object value, QueryShardContext context) { + return null; + } + }; + + when(mapperService.fieldType(any())).thenReturn(null); + + // Null fieldName + String fieldType = queryShardContext.getFieldTypeString(null); + assertEquals(null, fieldType); + + // MapperService return null + fieldType = queryShardContext.getFieldTypeString("field1"); + assertEquals(null, fieldType); + + when(mapperService.fieldType("field")).thenReturn(mappedFieldType); + + // Unmapped fieldName + fieldType = queryShardContext.getFieldTypeString("unknown_field"); + assertEquals(null, fieldType); + // Date fieldType + fieldType = queryShardContext.getFieldTypeString("field"); + assertEquals("long", fieldType); + } + } diff --git a/server/src/test/java/org/opensearch/index/query/plugin/DummyQueryBuilder.java b/server/src/test/java/org/opensearch/index/query/plugin/DummyQueryBuilder.java index 2b0aaa2dd7091..a32b8aaf7c2f7 100644 --- a/server/src/test/java/org/opensearch/index/query/plugin/DummyQueryBuilder.java +++ b/server/src/test/java/org/opensearch/index/query/plugin/DummyQueryBuilder.java @@ -62,6 +62,11 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep builder.startObject(NAME).endObject(); } + @Override + public final String fieldName() { + return getDefaultFieldName(); + } + public static DummyQueryBuilder fromXContent(XContentParser parser) throws IOException { XContentParser.Token token = parser.nextToken(); assert token == XContentParser.Token.END_OBJECT; diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java index e8a0f70ee3563..9fb4db96cbf7c 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java @@ -900,6 +900,11 @@ protected void doWriteTo(StreamOutput out) {} @Override protected void doXContent(XContentBuilder builder, Params params) {} + @Override + public String fieldName() { + return getDefaultFieldName(); + } + @Override protected Query doToQuery(QueryShardContext context) { return null; diff --git a/server/src/test/java/org/opensearch/search/sort/AbstractSortTestCase.java b/server/src/test/java/org/opensearch/search/sort/AbstractSortTestCase.java index 257ff1015e3b4..9634c80f26fb8 100644 --- a/server/src/test/java/org/opensearch/search/sort/AbstractSortTestCase.java +++ b/server/src/test/java/org/opensearch/search/sort/AbstractSortTestCase.java @@ -55,6 +55,7 @@ import org.opensearch.index.mapper.ContentPath; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.Mapper.BuilderContext; +import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.NumberFieldMapper; import org.opensearch.index.mapper.ObjectMapper; import org.opensearch.index.mapper.ObjectMapper.Nested; @@ -199,10 +200,18 @@ public void testEqualsAndHashcode() { } protected final QueryShardContext createMockShardContext() { - return createMockShardContext(null); + return createMockShardContext(null, null); } protected final QueryShardContext createMockShardContext(IndexSearcher searcher) { + return createMockShardContext(searcher, null); + } + + protected final QueryShardContext createMockShardContext(MapperService mockMapperService) { + return createMockShardContext(null, mockMapperService); + } + + protected final QueryShardContext createMockShardContext(IndexSearcher searcher, MapperService mapperService) { Index index = new Index(randomAlphaOfLengthBetween(1, 10), "_na_"); IndexSettings idxSettings = IndexSettingsModule.newIndexSettings( index, @@ -222,7 +231,7 @@ protected final QueryShardContext createMockShardContext(IndexSearcher searcher) BigArrays.NON_RECYCLING_INSTANCE, bitsetFilterCache, indexFieldDataLookup, - null, + mapperService, null, scriptService, xContentRegistry(), diff --git a/server/src/test/java/org/opensearch/search/sort/FieldSortBuilderTests.java b/server/src/test/java/org/opensearch/search/sort/FieldSortBuilderTests.java index 9b8cd1b5f1ce0..3a889fa0bc551 100644 --- a/server/src/test/java/org/opensearch/search/sort/FieldSortBuilderTests.java +++ b/server/src/test/java/org/opensearch/search/sort/FieldSortBuilderTests.java @@ -46,6 +46,7 @@ import org.apache.lucene.sandbox.document.BigIntegerPoint; import org.apache.lucene.sandbox.document.HalfFloatPoint; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; import org.apache.lucene.search.SortField; import org.apache.lucene.search.SortedNumericSelector; import org.apache.lucene.search.SortedNumericSortField; @@ -64,8 +65,11 @@ import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.NestedPathFieldMapper; import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.index.mapper.TextSearchInfo; +import org.opensearch.index.mapper.ValueFetcher; import org.opensearch.index.query.MatchNoneQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; @@ -77,11 +81,13 @@ import org.opensearch.search.MultiValueMode; import org.opensearch.search.SearchSortValuesAndFormats; import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.lookup.SearchLookup; import java.io.IOException; import java.math.BigInteger; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Locale; @@ -91,6 +97,8 @@ import static org.opensearch.search.sort.FieldSortBuilder.getPrimaryFieldSortOrNull; import static org.opensearch.search.sort.NestedSortBuilderTests.createRandomNestedSort; import static org.hamcrest.Matchers.instanceOf; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class FieldSortBuilderTests extends AbstractSortTestCase { @@ -704,6 +712,36 @@ public void testIsBottomSortShardDisjoint() throws Exception { } } + /** + * Test that the sort builder fieldType is set properly + */ + public void testSortFieldFieldType() throws IOException { + MapperService mapperService = mock(MapperService.class); + QueryShardContext shardContextMock = createMockShardContext(mapperService); + + MappedFieldType mappedFieldType = new MappedFieldType("field", true, false, true, TextSearchInfo.NONE, Collections.emptyMap()) { + @Override + public String typeName() { + return "double"; + } + + @Override + public ValueFetcher valueFetcher(QueryShardContext context, SearchLookup searchLookup, String format) { + throw new UnsupportedOperationException(); + } + + @Override + public Query termQuery(Object value, QueryShardContext context) { + return null; + } + }; + + when(mapperService.fieldType("field")).thenReturn(mappedFieldType); + FieldSortBuilder fieldSortBuilder = new FieldSortBuilder("field"); + fieldSortBuilder.buildHelper(shardContextMock); + assertEquals(fieldSortBuilder.getFieldType(), "double"); + } + @Override protected void assertWarnings(FieldSortBuilder testItem) { List expectedWarnings = new ArrayList<>(); diff --git a/server/src/test/java/org/opensearch/search/sort/plugin/CustomSortBuilder.java b/server/src/test/java/org/opensearch/search/sort/plugin/CustomSortBuilder.java index c3d790cb99cbd..56b7b9d86c253 100644 --- a/server/src/test/java/org/opensearch/search/sort/plugin/CustomSortBuilder.java +++ b/server/src/test/java/org/opensearch/search/sort/plugin/CustomSortBuilder.java @@ -75,6 +75,11 @@ public BucketedSort buildBucketedSort(final QueryShardContext context, final int throw new IllegalStateException("rewrite"); } + @Override + public String fieldName() { + return field; + } + @Override public boolean equals(Object object) { if (this == object) {