A bit of cleanup

Author: kderusso

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java

@@ -58,7 +58,6 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.stream.Collectors;
 
 /**
@@ -307,24 +306,26 @@ public void onFailure(Exception exc) {
                 return;
             }
             int currentBatchSize = Math.min(requests.size(), batchSize);
-            final List<FieldInferenceRequest> currentBatch = requests.subList(0, currentBatchSize);
+            ChunkingSettings chunkingSettings = requests.get(0).chunkingSettings;
+            List<FieldInferenceRequest> currentBatch = new ArrayList<>();
+            List<FieldInferenceRequest> others = new ArrayList<>();
+            for (int i = 0; i < currentBatchSize; i++) {
+                FieldInferenceRequest request = requests.get(i);
+                if ((chunkingSettings == null && request.chunkingSettings == null) || request.chunkingSettings.equals(chunkingSettings)) {
+                    currentBatch.add(request);
+                } else {
+                    others.add(request);
+                }
+            }
+
             final List<FieldInferenceRequest> nextBatch = requests.subList(currentBatchSize, requests.size());
-            //
-            // List<ChunkedInputs> chunkedInputs = currentBatch.stream()
-            // .map(request -> new ChunkedInputs(request.chunkingSettings(), List.of(request.input())))
-            // .toList();
-
-            List<ChunkedInputs> chunkedInputs = currentBatch.stream()
-                .collect(Collectors.groupingBy(request -> Optional.ofNullable(request.chunkingSettings())))
-                .entrySet()
-                .stream()
-                .map(
-                    entry -> new ChunkedInputs(
-                        entry.getKey().orElse(null),
-                        entry.getValue().stream().map(FieldInferenceRequest::input).collect(Collectors.toList())
-                    )
-                )
-                .toList();
+            nextBatch.addAll(others);
+
+            // We can assume current batch has all the same chunking settings
+            ChunkedInputs chunkedInputs = new ChunkedInputs(
+                chunkingSettings,
+                currentBatch.stream().map(r -> r.input).collect(Collectors.toList())
+            );
 
             ActionListener<List<ChunkedInference>> completionListener = new ActionListener<>() {
                 @Override
@@ -390,19 +391,17 @@ private void onFinish() {
                 }
             };
 
-            for (ChunkedInputs chunkedInput : chunkedInputs) {
-                inferenceProvider.service()
-                    .chunkedInfer(
-                        inferenceProvider.model(),
-                        null,
-                        chunkedInput.inputs(),
-                        Map.of(),
-                        chunkedInput.chunkingSettings(),
-                        InputType.INGEST,
-                        TimeValue.MAX_VALUE,
-                        completionListener
-                    );
-            }
+            inferenceProvider.service()
+                .chunkedInfer(
+                    inferenceProvider.model(),
+                    null,
+                    chunkedInputs.inputs(),
+                    Map.of(),
+                    chunkedInputs.chunkingSettings(),
+                    InputType.INGEST,
+                    TimeValue.MAX_VALUE,
+                    completionListener
+                );
         }
 
         private FieldInferenceResponseAccumulator ensureResponseAccumulatorSlot(int id) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldTests.java

@@ -70,6 +70,7 @@ protected void assertEqualInstances(SemanticTextField expectedInstance, Semantic
         assertThat(newInstance.originalValues(), equalTo(expectedInstance.originalValues()));
         assertThat(newInstance.inference().modelSettings(), equalTo(expectedInstance.inference().modelSettings()));
         assertThat(newInstance.inference().chunks().size(), equalTo(expectedInstance.inference().chunks().size()));
+        assertThat(newInstance.chunkingSettings(), equalTo(expectedInstance.chunkingSettings()));
         MinimalServiceSettings modelSettings = newInstance.inference().modelSettings();
         for (var entry : newInstance.inference().chunks().entrySet()) {
             var expectedChunks = expectedInstance.inference().chunks().get(entry.getKey());