Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add script fields support for multi search request #632

Merged
merged 4 commits into from
Sep 30, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)

### Added
- Added support for "smartcn" analyzer ([#605](https://github.com/opensearch-project/opensearch-java/pull/605))
- Added support for "script_fields" in multi search request ([#632](https://github.com/opensearch-project/opensearch-java/pull/632))

### Dependencies

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@

package org.opensearch.client.opensearch.core.msearch;

import org.opensearch.client.opensearch._types.ScriptField;
import org.opensearch.client.opensearch._types.SortOptions;
import org.opensearch.client.opensearch._types.aggregations.Aggregation;
import org.opensearch.client.opensearch._types.query_dsl.Query;
Expand Down Expand Up @@ -96,6 +97,8 @@ public class MultisearchBody implements JsonpSerializable {
@Nullable
private SourceConfig source;

private final Map<String, ScriptField> scriptFields;

// ---------------------------------------------------------------------------------------------

private MultisearchBody(Builder builder) {
Expand All @@ -113,7 +116,7 @@ private MultisearchBody(Builder builder) {
this.suggest = builder.suggest;
this.highlight = builder.highlight;
this.source = builder.source;

this.scriptFields = ApiTypeHelper.unmodifiable(builder.scriptFields);
}

public static MultisearchBody of(Function<Builder, ObjectBuilder<MultisearchBody>> fn) {
Expand Down Expand Up @@ -221,6 +224,10 @@ public final SourceConfig source() {
return this.source;
}

public final Map<String, ScriptField> scriptFields() {
return this.scriptFields;
}

/**
* Serialize this object to JSON.
*/
Expand Down Expand Up @@ -316,6 +323,18 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {

}

if (ApiTypeHelper.isDefined(this.scriptFields)) {
generator.writeKey("script_fields");
generator.writeStartObject();
for (Map.Entry<String, ScriptField> item0 : this.scriptFields.entrySet()) {
generator.writeKey(item0.getKey());
item0.getValue().serialize(generator, mapper);

}
generator.writeEnd();

}

}

// ---------------------------------------------------------------------------------------------
Expand Down Expand Up @@ -364,6 +383,8 @@ public static class Builder extends ObjectBuilderBase implements ObjectBuilder<M
@Nullable
private SourceConfig source;

private Map<String, ScriptField> scriptFields;
reta marked this conversation as resolved.
Show resolved Hide resolved


/**
* API name: {@code aggregations}
Expand Down Expand Up @@ -571,6 +592,29 @@ public final Builder source(Function<SourceConfig.Builder, ObjectBuilder<SourceC
return this.source(fn.apply(new SourceConfig.Builder()).build());
}

/**
* API name: {@code script_fields}
*/
public final Builder scriptFields(Map<String, ScriptField> map) {
this.scriptFields = _mapPutAll(this.scriptFields, map);
return this;
}

/**
* API name: {@code script_fields}
*/
public final Builder scriptFields(String key, ScriptField value) {
this.scriptFields = _mapPut(this.scriptFields, key, value);
return this;
}

/**
* API name: {@code script_fields}
*/
public final Builder scriptFields(String key, Function<ScriptField.Builder, ObjectBuilder<ScriptField>> fn) {
return scriptFields(key, fn.apply(new ScriptField.Builder()).build());
}

/**
* Builds a {@link MultisearchBody}.
*
Expand Down Expand Up @@ -609,6 +653,8 @@ protected static void setupMultisearchBodyDeserializer(ObjectDeserializer<Multis
op.add(Builder::suggest, Suggester._DESERIALIZER, "suggest");
op.add(Builder::highlight, Highlight._DESERIALIZER, "highlight");
op.add(Builder::source, SourceConfig._DESERIALIZER, "_source");
op.add(Builder::scriptFields, JsonpDeserializer.stringMapDeserializer(ScriptField._DESERIALIZER),
"script_fields");

}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,13 @@
package org.opensearch.client.opensearch.integTest;

import org.junit.Test;
import org.opensearch.client.json.JsonData;
import org.opensearch.client.opensearch._types.FieldSort;
import org.opensearch.client.opensearch._types.FieldValue;
import org.opensearch.client.opensearch._types.InlineScript;
import org.opensearch.client.opensearch._types.Refresh;
import org.opensearch.client.opensearch._types.Script;
import org.opensearch.client.opensearch._types.ScriptField;
import org.opensearch.client.opensearch._types.SortOptions;
import org.opensearch.client.opensearch._types.SortOrder;
import org.opensearch.client.opensearch._types.query_dsl.FuzzyQuery;
Expand All @@ -27,8 +31,13 @@
import org.opensearch.client.opensearch.core.search.SourceConfig;
import org.opensearch.client.util.ObjectBuilder;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

public abstract class AbstractMultiSearchRequestIT extends OpenSearchJavaClientTestCase {
Expand Down Expand Up @@ -175,6 +184,44 @@ public void shouldReturnMultiSearchesSearchAfter() throws Exception {
assertEquals(1, response2.responses().get(0).result().hits().hits().size());
}

@Test
public void shouldReturnMultiSearchesScriptFields() throws Exception {
String index = "multiple_searches_request_script_fields";
createTestDocuments(index);

RequestItem sortedItemsQuery = createMSearchSortedFuzzyRequest();

MsearchResponse<ShopItem> response = sendMSearchRequest(index, List.of(sortedItemsQuery));
assertEquals(1, response.responses().size());
var hits = response.responses().get(0).result().hits().hits();
assertEquals(3, hits.size());
assertNull(hits.get(0).score());
assertNull(hits.get(1).score());
assertNull(hits.get(2).score());

Map<String, ScriptField> scriptFields = new HashMap<>();
scriptFields.put("quantity", new ScriptField.Builder().script(Script.of(s -> s.inline(new InlineScript.Builder()
.lang("painless")
.source("doc['quantity'].value + params.inc")
.params("inc", JsonData.of(1))
.build())))
.build());


RequestItem requestItem = createMSearchQueryWithScriptFields("small", scriptFields);

MsearchResponse<ShopItem> responseWithScriptFields = sendMSearchRequest(index, List.of(requestItem));
var hitsWithScriptFields = responseWithScriptFields.responses().get(0).result().hits().hits();
assertEquals(2, hitsWithScriptFields.size());
// validating that the quantity for small items is increased by 1
ObjectMapper mapper = new ObjectMapper();
JsonNode node = mapper.readTree(hitsWithScriptFields.get(0).fields().get("quantity").toString());
assertEquals(2, (int) mapper.treeToValue(node.get(0), int.class));

node = mapper.readTree(hitsWithScriptFields.get(1).fields().get("quantity").toString());
assertEquals(3, (int) mapper.treeToValue(node.get(0), int.class));
}


private void assertResponseSources(MultiSearchResponseItem<ShopItem> response) {
List<Hit<ShopItem>> hitsWithHighlights = response.result().hits().hits();
Expand All @@ -198,6 +245,14 @@ private RequestItem createMSearchQuery(String itemSize) {
return createMSearchQuery(itemSize, null, List.of());
}

private RequestItem createMSearchQueryWithScriptFields(String itemSize, Map<String, ScriptField> scriptFields) {
return RequestItem.of(item -> item.header(header -> header)
.body(body -> body.query(createItemSizeSearchQuery(itemSize))
.scriptFields(scriptFields)
)
);
}

private RequestItem createMSearchQueryWithHighlight(String itemSize) {
return createMSearchQuery(itemSize, "size", List.of());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
* Run with: <c>./gradlew :samples:run -Dsamples.mainClass=Search</c>
*/
public class Search {
private static final Logger LOGGER = LogManager.getLogger(IndexingBasics.class);
private static final Logger LOGGER = LogManager.getLogger(Search.class);

private static OpenSearchClient client;

Expand Down
Loading