[LTR] Fix feature display order when using explain. (#137671) (#137684)

Author: afoucret

docs/changelog/137671.yaml

@@ -0,0 +1,5 @@
+pr: 137671
+summary: "[LTR] Fix feature display order when using explain"
+area: Search
+type: bug
+issues: []

x-pack/plugin/ml/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/LearningToRankRescorerIT.java

@@ -83,7 +83,7 @@ public void testLearningToRankRescoreWithExplain() throws Exception {
               }
             }""");
         var response = client().performRequest(request);
-        assertExplainExtractedFeatures(response, List.of("type_tv", "cost", "two"));
+        assertExplainExtractedFeatures(response, List.of("cost", "type_tv", "two"));
     }
 
     public void testLearningToRankRescoreSmallWindow() throws Exception {
@@ -192,27 +192,25 @@ private static void assertExplainExtractedFeatures(Response response, List<Strin
             assertThat(queryDetails.get(1).get("description"), equalTo("extracted features"));
 
             var featureDetails = new ArrayList<>((ArrayList<Map<String, Object>>) queryDetails.get(1).get("details"));
-            assertThat(featureDetails.size(), equalTo(3));
-
-            var missingKeys = new ArrayList<String>();
-            for (String expectedFeature : expectedFeatures) {
-                var expectedDescription = Strings.format("feature value for [%s]", expectedFeature);
-
-                var wasFound = false;
-                for (Map<String, Object> detailItem : featureDetails) {
-                    if (detailItem.get("description").equals(expectedDescription)) {
-                        featureDetails.remove(detailItem);
-                        wasFound = true;
-                        break;
-                    }
-                }
-
-                if (wasFound == false) {
-                    missingKeys.add(expectedFeature);
+            assertThat(featureDetails.size(), equalTo(expectedFeatures.size()));
+
+            // Extract feature names in the order they appear in the explanation
+            List<String> actualFeatureOrder = new ArrayList<>();
+            for (Map<String, Object> detailItem : featureDetails) {
+                String description = (String) detailItem.get("description");
+                // Extract feature name from "feature value for [featureName]"
+                if (description != null && description.startsWith("feature value for [") && description.endsWith("]")) {
+                    String featureName = description.substring("feature value for [".length(), description.length() - 1);
+                    actualFeatureOrder.add(featureName);
                 }
             }
 
-            assertThat(Strings.format("Could not find features: [%s]", String.join(", ", missingKeys)), featureDetails.size(), equalTo(0));
+            // Verify that features appear in the expected order
+            assertThat(
+                "Features should appear in the expected order. Expected: " + expectedFeatures + ", Actual: " + actualFeatureOrder,
+                actualFeatureOrder,
+                equalTo(expectedFeatures)
+            );
         }
     }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/ltr/LearningToRankRescorer.java

@@ -28,7 +28,6 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 
 import static java.util.stream.Collectors.toUnmodifiableSet;
@@ -164,21 +163,24 @@ public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreCon
         List<FeatureExtractor> featureExtractors = ltrContext.buildFeatureExtractors(searcher);
         int featureSize = featureExtractors.stream().mapToInt(fe -> fe.featureNames().size()).sum();
 
-        Map<String, Object> features = Maps.newMapWithExpectedSize(featureSize);
+        Map<String, Object> extractedFeatures = Maps.newMapWithExpectedSize(featureSize);
 
         for (FeatureExtractor featureExtractor : featureExtractors) {
             featureExtractor.setNextReader(currentSegment);
-            featureExtractor.addFeatures(features, targetDoc);
+            featureExtractor.addFeatures(extractedFeatures, targetDoc);
         }
 
         // Predicting the value
-        var ltrScore = ((Number) localModelDefinition.inferLtr(features, ltrContext.learningToRankConfig).predictedValue()).floatValue();
+        var ltrScore = ((Number) localModelDefinition.inferLtr(extractedFeatures, ltrContext.learningToRankConfig).predictedValue())
+            .floatValue();
 
         List<Explanation> featureExplanations = new ArrayList<>();
-        for (String featureName : features.keySet()) {
-            Number featureValue = Objects.requireNonNullElse((Number) features.get(featureName), 0);
-            featureExplanations.add(Explanation.match(featureValue, "feature value for [" + featureName + "]"));
-        }
+        ltrContext.learningToRankConfig.getFeatureExtractorBuilders().forEach(featureExtractor -> {
+            String featureName = featureExtractor.featureName();
+            if (extractedFeatures.containsKey(featureName) && extractedFeatures.get(featureName) instanceof Number featureValue) {
+                featureExplanations.add(Explanation.match(featureValue, "feature value for [" + featureName + "]"));
+            }
+        });
 
         return Explanation.match(
             ltrScore,