use stream method

Signed-off-by: Mingshi Liu <[email protected]>
opensearch-project · Oct 16, 2024 · ccf9c61 · ccf9c61
1 parent 8381354
commit ccf9c61
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchResponseProcessor.java b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchResponseProcessor.java
@@ -804,18 +804,15 @@ public MLInferenceSearchResponseProcessor create(
             }
             boolean writeToSearchExtension = false;
 
-            if (outputMaps != null) {
-                for (Map<String, String> outputMap : outputMaps) {
-                    for (String key : outputMap.keySet()) {
-                        if (key.startsWith(EXTENSION_PREFIX)) {
-                            writeToSearchExtension = true;
-                            break;
-                        }
-                    }
-                }
+            if (outputMaps != null
+                && outputMaps
+                    .stream()
+                    .anyMatch(outputMap -> outputMap.keySet().stream().anyMatch(key -> key.startsWith(EXTENSION_PREFIX)))) {
+                writeToSearchExtension = true;
             }
+
             if (writeToSearchExtension & oneToOne) {
-                throw new IllegalArgumentException("Writing model response to search extension does not support when one_to_one is true.");
+                throw new IllegalArgumentException("Write model response to search extension does not support when one_to_one is true.");
             }
 
             return new MLInferenceSearchResponseProcessor(