Add script fields support for multi search request (#632) (#646)

* Adding script_fields support for mseearch request with tests * Fixing logger in Search sample * Fixing build * Fixing tests --------- Signed-off-by: Vacha Shah <[email protected]>
opensearch-project · Oct 3, 2023 · d7b3abf · d7b3abf
1 parent d498c65
commit d7b3abf
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Added support for "smartcn" analyzer ([#605](https://github.com/opensearch-project/opensearch-java/pull/605))
 - Added support for "cjk" analyzer ([#604](https://github.com/opensearch-project/opensearch-java/pull/604))
 - Added support for wrapper queries ([#630](https://github.com/opensearch-project/opensearch-java/pull/630))
+- Added support for "script_fields" in multi search request ([#632](https://github.com/opensearch-project/opensearch-java/pull/632))
 
 ### Dependencies
 - Bumps `org.ajoberstar.grgit:grgit-gradle` from 5.0.0 to 5.2.0

diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java b/java-client/src/main/java/org/opensearch/client/opensearch/core/msearch/MultisearchBody.java
@@ -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;
@@ -96,6 +97,8 @@ public class MultisearchBody implements JsonpSerializable {
 	@Nullable
 	private SourceConfig source;
 
+	private final Map<String, ScriptField> scriptFields;
+
 	// ---------------------------------------------------------------------------------------------
 
 	private MultisearchBody(Builder builder) {
@@ -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) {
@@ -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.
 	 */
@@ -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();
+
+		}
+
 	}
 
 	// ---------------------------------------------------------------------------------------------
@@ -364,6 +383,8 @@ public static class Builder extends ObjectBuilderBase implements ObjectBuilder<M
 		@Nullable
 		private SourceConfig source;
 
+		private Map<String, ScriptField> scriptFields;
+
 
 		/**
 		 * API name: {@code aggregations}
@@ -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}.
 		 *
@@ -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");
 
 	}
 

diff --git a/...rc/test/java/org/opensearch/client/opensearch/integTest/AbstractMultiSearchRequestIT.java b/...rc/test/java/org/opensearch/client/opensearch/integTest/AbstractMultiSearchRequestIT.java
@@ -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;
@@ -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 {
@@ -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();
@@ -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());
 	}

diff --git a/samples/src/main/java/org/opensearch/client/samples/Search.java b/samples/src/main/java/org/opensearch/client/samples/Search.java
@@ -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;