Check for inference errors on the coordinator node

Author: Mikep86

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java

@@ -238,8 +238,8 @@ private QueryBuilder doRewriteBuildSemanticQuery(SearchExecutionContext searchEx
                 inferenceResults = inferenceResultsMap.get(inferenceId);
             }
 
-            return switch (inferenceResults) {
-                case null -> throw new IllegalStateException(
+            if (inferenceResults == null) {
+                throw new IllegalStateException(
                     "No inference results set for ["
                         + semanticTextFieldType.typeName()
                         + "] field ["
@@ -248,25 +248,9 @@ private QueryBuilder doRewriteBuildSemanticQuery(SearchExecutionContext searchEx
                         + inferenceId
                         + "]"
                 );
-                case ErrorInferenceResults errorInferenceResults -> throw new InferenceException(
-                    "Field [" + fieldName + "] with inference ID [" + inferenceId + "] query inference error",
-                    errorInferenceResults.getException()
-                );  // Use InferenceException here so that the status code is set by the cause
-                case WarningInferenceResults warningInferenceResults -> throw new IllegalStateException(
-                    "Field ["
-                        + fieldName
-                        + "] with inference ID ["
-                        + inferenceId
-                        + "] query inference warning: "
-                        + warningInferenceResults.getWarning()
-                );
-                default -> semanticTextFieldType.semanticQuery(
-                    inferenceResults,
-                    searchExecutionContext.requestSize(),
-                    boost(),
-                    queryName()
-                );
-            };
+            }
+
+            return semanticTextFieldType.semanticQuery(inferenceResults, searchExecutionContext.requestSize(), boost(), queryName());
         } else if (lenient != null && lenient) {
             return new MatchNoneQueryBuilder();
         } else {
@@ -278,6 +262,7 @@ private QueryBuilder doRewriteBuildSemanticQuery(SearchExecutionContext searchEx
 
     private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext queryRewriteContext) {
         if (inferenceResultsMap != null) {
+            inferenceResultsErrorCheck();
             return this;
         }
 
@@ -378,6 +363,30 @@ private static InferenceResults validateAndConvertInferenceResults(
         return inferenceResults;
     }
 
+    private void inferenceResultsErrorCheck() {
+        for (var entry : inferenceResultsMap.entrySet()) {
+            String inferenceId = entry.getKey();
+            InferenceResults inferenceResults = entry.getValue();
+
+            if (inferenceResults instanceof ErrorInferenceResults errorInferenceResults) {
+                // Use InferenceException here so that the status code is set by the cause
+                throw new InferenceException(
+                    "Field [" + fieldName + "] with inference ID [" + inferenceId + "] query inference error",
+                    errorInferenceResults.getException()
+                );
+            } else if (inferenceResults instanceof WarningInferenceResults warningInferenceResults) {
+                throw new IllegalStateException(
+                    "Field ["
+                        + fieldName
+                        + "] with inference ID ["
+                        + inferenceId
+                        + "] query inference warning: "
+                        + warningInferenceResults.getWarning()
+                );
+            }
+        }
+    }
+
     @Override
     protected Query doToQuery(SearchExecutionContext context) throws IOException {
         throw new IllegalStateException(NAME + " should have been rewritten to another query type");