Skip to content

Commit

Permalink
Add SearchExtBuilders to SearchResponse (#9379)
Browse files Browse the repository at this point in the history
* Add SearchExtBuilders to SearchResponse. [Issue #9328](#9328)

Signed-off-by: Austin Lee <[email protected]>

* Keep SearchResponse immutable, add a constructor to take a List of SearchExtBuilders.

Signed-off-by: Austin Lee <[email protected]>

* Fix spotlessJavaCheck findings.

Signed-off-by: Austin Lee <[email protected]>

* Move SearchExtBuilders into SearchResponseSections, fix indenting in SearchRequest.

Signed-off-by: Austin Lee <[email protected]>

* Updated changelog (mixed minor formatting issues), added version checks on serialization/deserialization, added a Builder for making copies of SearchResponse easier.

Signed-off-by: Austin Lee <[email protected]>

* Add GenericSearchExtBuilder as a catch-all for SearchExtBuilders not registered in xcontent registry.

Signed-off-by: Austin Lee <[email protected]>

* Simplify GenericSearchExtBuilder using a single Object member.

Signed-off-by: Austin Lee <[email protected]>

* Address additional review comments.

Signed-off-by: Austin Lee <[email protected]>

* Add Javadocs.

Signed-off-by: Austin Lee <[email protected]>

---------

Signed-off-by: Austin Lee <[email protected]>
  • Loading branch information
austintlee authored Aug 29, 2023
1 parent bb7d23c commit 60787b8
Show file tree
Hide file tree
Showing 7 changed files with 893 additions and 16 deletions.
11 changes: 6 additions & 5 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -78,17 +78,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
## [Unreleased 2.x]
### Added
- Add server version as REST response header [#6583](https:/opensearch-project/OpenSearch/issues/6583)
- Start replication checkpointTimers on primary before segments upload to remote store. ([#8221]()https:/opensearch-project/OpenSearch/pull/8221)
- [distribution/archives] [Linux] [x64] Provide the variant of the distributions bundled with JRE ([#8195]()https:/opensearch-project/OpenSearch/pull/8195)
- Start replication checkpointTimers on primary before segments upload to remote store. ([#8221](https:/opensearch-project/OpenSearch/pull/8221))
- [distribution/archives] [Linux] [x64] Provide the variant of the distributions bundled with JRE ([#8195](https:/opensearch-project/OpenSearch/pull/8195))
- Add configuration for file cache size to max remote data ratio to prevent oversubscription of file cache ([#8606](https:/opensearch-project/OpenSearch/pull/8606))
- Disallow compression level to be set for default and best_compression index codecs ([#8737]()https:/opensearch-project/OpenSearch/pull/8737)
- Disallow compression level to be set for default and best_compression index codecs ([#8737](https:/opensearch-project/OpenSearch/pull/8737))
- Prioritize replica shard movement during shard relocation ([#8875](https:/opensearch-project/OpenSearch/pull/8875))
- Introducing Default and Best Compression codecs as their algorithm name ([#9123]()https:/opensearch-project/OpenSearch/pull/9123)
- Make SearchTemplateRequest implement IndicesRequest.Replaceable ([#9122]()https:/opensearch-project/OpenSearch/pull/9122)
- Introducing Default and Best Compression codecs as their algorithm name ([#9123](https:/opensearch-project/OpenSearch/pull/9123))
- Make SearchTemplateRequest implement IndicesRequest.Replaceable ([#9122](https:/opensearch-project/OpenSearch/pull/9122))
- [BWC and API enforcement] Define the initial set of annotations, their meaning and relations between them ([#9223](https:/opensearch-project/OpenSearch/pull/9223))
- [Segment Replication] Support realtime reads for GET requests ([#9212](https:/opensearch-project/OpenSearch/pull/9212))
- [Feature] Expose term frequency in Painless script score context ([#9081](https:/opensearch-project/OpenSearch/pull/9081))
- Add support for reading partial files to HDFS repository ([#9513](https:/opensearch-project/OpenSearch/issues/9513))
- Add support for extensions to search responses using SearchExtBuilder ([#9379](https:/opensearch-project/OpenSearch/pull/9379))

### Dependencies
- Bump `org.apache.logging.log4j:log4j-core` from 2.17.1 to 2.20.0 ([#8307](https:/opensearch-project/OpenSearch/pull/8307))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,12 @@
import org.opensearch.core.xcontent.MediaTypeRegistry;
import org.opensearch.core.xcontent.ToXContentFragment;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParseException;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.core.xcontent.XContentParser.Token;
import org.opensearch.rest.action.RestActions;
import org.opensearch.search.GenericSearchExtBuilder;
import org.opensearch.search.SearchExtBuilder;
import org.opensearch.search.SearchHit;
import org.opensearch.search.SearchHits;
import org.opensearch.search.aggregations.Aggregations;
Expand All @@ -65,6 +68,7 @@
import java.util.Objects;
import java.util.function.Supplier;

import static org.opensearch.action.search.SearchResponseSections.EXT_FIELD;
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;

/**
Expand Down Expand Up @@ -312,6 +316,7 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
);
clusters.toXContent(builder, params);
internalResponse.toXContent(builder, params);

return builder;
}

Expand Down Expand Up @@ -339,6 +344,7 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
String searchContextId = null;
List<ShardSearchFailure> failures = new ArrayList<>();
Clusters clusters = Clusters.EMPTY;
List<SearchExtBuilder> extBuilders = new ArrayList<>();
for (Token token = parser.nextToken(); token != Token.END_OBJECT; token = parser.nextToken()) {
if (token == Token.FIELD_NAME) {
currentFieldName = parser.currentName();
Expand Down Expand Up @@ -417,6 +423,33 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
}
}
clusters = new Clusters(total, successful, skipped);
} else if (EXT_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
String extSectionName = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
extSectionName = parser.currentName();
} else {
SearchExtBuilder searchExtBuilder;
try {
searchExtBuilder = parser.namedObject(SearchExtBuilder.class, extSectionName, null);
if (!searchExtBuilder.getWriteableName().equals(extSectionName)) {
throw new IllegalStateException(
"The parsed ["
+ searchExtBuilder.getClass().getName()
+ "] object has a "
+ "different writeable name compared to the name of the section that it was parsed from: found ["
+ searchExtBuilder.getWriteableName()
+ "] expected ["
+ extSectionName
+ "]"
);
}
} catch (XContentParseException e) {
searchExtBuilder = GenericSearchExtBuilder.fromXContent(parser);
}
extBuilders.add(searchExtBuilder);
}
}
} else {
parser.skipChildren();
}
Expand All @@ -429,7 +462,8 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
timedOut,
terminatedEarly,
profile,
numReducePhases
numReducePhases,
extBuilders
);
return new SearchResponse(
searchResponseSections,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,18 +32,23 @@

package org.opensearch.action.search;

import org.opensearch.core.ParseField;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.ToXContentFragment;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.search.SearchExtBuilder;
import org.opensearch.search.SearchHits;
import org.opensearch.search.aggregations.Aggregations;
import org.opensearch.search.profile.ProfileShardResult;
import org.opensearch.search.profile.SearchProfileShardResults;
import org.opensearch.search.suggest.Suggest;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
* Base class that holds the various sections which a search response is
Expand All @@ -57,13 +62,16 @@
*/
public class SearchResponseSections implements ToXContentFragment {

public static final ParseField EXT_FIELD = new ParseField("ext");

protected final SearchHits hits;
protected final Aggregations aggregations;
protected final Suggest suggest;
protected final SearchProfileShardResults profileResults;
protected final boolean timedOut;
protected final Boolean terminatedEarly;
protected final int numReducePhases;
protected final List<SearchExtBuilder> searchExtBuilders = new ArrayList<>();

public SearchResponseSections(
SearchHits hits,
Expand All @@ -73,6 +81,19 @@ public SearchResponseSections(
Boolean terminatedEarly,
SearchProfileShardResults profileResults,
int numReducePhases
) {
this(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, Collections.emptyList());
}

public SearchResponseSections(
SearchHits hits,
Aggregations aggregations,
Suggest suggest,
boolean timedOut,
Boolean terminatedEarly,
SearchProfileShardResults profileResults,
int numReducePhases,
List<SearchExtBuilder> searchExtBuilders
) {
this.hits = hits;
this.aggregations = aggregations;
Expand All @@ -81,6 +102,7 @@ public SearchResponseSections(
this.timedOut = timedOut;
this.terminatedEarly = terminatedEarly;
this.numReducePhases = numReducePhases;
this.searchExtBuilders.addAll(Objects.requireNonNull(searchExtBuilders, "searchExtBuilders must not be null"));
}

public final boolean timedOut() {
Expand Down Expand Up @@ -135,9 +157,20 @@ public final XContentBuilder toXContent(XContentBuilder builder, Params params)
if (profileResults != null) {
profileResults.toXContent(builder, params);
}
if (!searchExtBuilders.isEmpty()) {
builder.startObject(EXT_FIELD.getPreferredName());
for (SearchExtBuilder searchExtBuilder : searchExtBuilders) {
searchExtBuilder.toXContent(builder, params);
}
builder.endObject();
}
return builder;
}

public List<SearchExtBuilder> getSearchExtBuilders() {
return Collections.unmodifiableList(this.searchExtBuilders);
}

protected void writeTo(StreamOutput out) throws IOException {
throw new UnsupportedOperationException();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,165 @@
/*
* 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;

import org.opensearch.core.ParseField;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParseException;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
* This is a catch-all SearchExtBuilder implementation that is used when an appropriate SearchExtBuilder
* is not found during SearchResponse's fromXContent operation.
*/
public final class GenericSearchExtBuilder extends SearchExtBuilder {

public final static ParseField EXT_BUILDER_NAME = new ParseField("generic_ext");

private final Object genericObj;
private final ValueType valueType;

enum ValueType {
SIMPLE(0),
MAP(1),
LIST(2);

private final int value;

ValueType(int value) {
this.value = value;
}

public int getValue() {
return value;
}

static ValueType fromInt(int value) {
switch (value) {
case 0:
return SIMPLE;
case 1:
return MAP;
case 2:
return LIST;
default:
throw new IllegalArgumentException("Unsupported value: " + value);
}
}
}

public GenericSearchExtBuilder(Object genericObj, ValueType valueType) {
this.genericObj = genericObj;
this.valueType = valueType;
}

public GenericSearchExtBuilder(StreamInput in) throws IOException {
valueType = ValueType.fromInt(in.readInt());
switch (valueType) {
case SIMPLE:
genericObj = in.readGenericValue();
break;
case MAP:
genericObj = in.readMap();
break;
case LIST:
genericObj = in.readList(r -> r.readGenericValue());
break;
default:
throw new IllegalStateException("Unable to construct GenericSearchExtBuilder from incoming stream.");
}
}

public static GenericSearchExtBuilder fromXContent(XContentParser parser) throws IOException {
// Look at the parser's next token.
// If it's START_OBJECT, parse as map, if it's START_ARRAY, parse as list, else
// parse as simpleVal
XContentParser.Token token = parser.currentToken();
ValueType valueType;
Object genericObj;
if (token == XContentParser.Token.START_OBJECT) {
genericObj = parser.map();
valueType = ValueType.MAP;
} else if (token == XContentParser.Token.START_ARRAY) {
genericObj = parser.list();
valueType = ValueType.LIST;
} else if (token.isValue()) {
genericObj = parser.objectText();
valueType = ValueType.SIMPLE;
} else {
throw new XContentParseException("Unknown token: " + token);
}

return new GenericSearchExtBuilder(genericObj, valueType);
}

@Override
public String getWriteableName() {
return EXT_BUILDER_NAME.getPreferredName();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeInt(valueType.getValue());
switch (valueType) {
case SIMPLE:
out.writeGenericValue(genericObj);
break;
case MAP:
out.writeMap((Map<String, Object>) genericObj);
break;
case LIST:
out.writeCollection((List<Object>) genericObj, StreamOutput::writeGenericValue);
break;
default:
throw new IllegalStateException("Unknown valueType: " + valueType);
}
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
switch (valueType) {
case SIMPLE:
return builder.field(EXT_BUILDER_NAME.getPreferredName(), genericObj);
case MAP:
return builder.field(EXT_BUILDER_NAME.getPreferredName(), (Map<String, Object>) genericObj);
case LIST:
return builder.field(EXT_BUILDER_NAME.getPreferredName(), (List<Object>) genericObj);
default:
return null;
}
}

// We need this for the equals method.
Object getValue() {
return genericObj;
}

@Override
public int hashCode() {
return Objects.hash(this.valueType, this.genericObj);
}

@Override
public boolean equals(Object obj) {
if (obj == null) {
return false;
}
if (!(obj instanceof GenericSearchExtBuilder)) {
return false;
}
return Objects.equals(getValue(), ((GenericSearchExtBuilder) obj).getValue())
&& Objects.equals(valueType, ((GenericSearchExtBuilder) obj).valueType);
}
}
Loading

0 comments on commit 60787b8

Please sign in to comment.