Add include_individual_nodes, include_all_nodes, include_info params to stats API

CHANGELOG.md

@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Stats] Add stats tracking for semantic field ([#1362](https://github.com/opensearch-project/neural-search/pull/1362))
 - [Stats] Add stats for neural query enricher, neural sparse encoding, two phase, and reranker processors ([#1343](https://github.com/opensearch-project/neural-search/pull/1343))
 
+- [Stats] Add `include_individual_nodes`, `include_all_nodes`, `include_info` parameters to stats API ([#1360](https://github.com/opensearch-project/neural-search/pull/1360))
 ### Bug Fixes
 - Fix score value as null for single shard when sorting is not done on score field ([#1277](https://github.com/opensearch-project/neural-search/pull/1277))
 - Return bad request for stats API calls with invalid stat names instead of ignoring them ([#1291](https://github.com/opensearch-project/neural-search/pull/1291))

qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/RestNeuralStatsActionIT.java

@@ -119,4 +119,3 @@ public void testStats_E2EFlow() throws Exception {
         }
     }
 }
-

src/main/java/org/opensearch/neuralsearch/common/MinClusterVersionUtil.java

@@ -27,6 +27,7 @@ public final class MinClusterVersionUtil {
     private static final Version MINIMAL_SUPPORTED_VERSION_PAGINATION_IN_HYBRID_QUERY = Version.V_2_19_0;
     private static final Version MINIMAL_SUPPORTED_VERSION_NEURAL_ORIGINAL_QUERY_TEXT = Version.V_3_0_0;
     public static final Version MINIMAL_SUPPORTED_VERSION_SEMANTIC_FIELD = Version.V_3_1_0;
+    public static final Version MINIMAL_SUPPORTED_VERSION_STATS_CATEGORY_FILTERING = Version.V_3_1_0;
 
     // Note this minimal version will act as an override
     private static final Map<String, Version> MINIMAL_VERSION_NEURAL = ImmutableMap.<String, Version>builder()
@@ -48,6 +49,10 @@ public static boolean isClusterOnOrAfterMinReqVersionForPaginationInHybridQuery(
         return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_PAGINATION_IN_HYBRID_QUERY);
     }
 
+    public static boolean isClusterOnOrAfterMinReqVersionForStatCategoryFiltering() {
+        return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_STATS_CATEGORY_FILTERING);
+    }
+
     public static boolean isClusterOnOrAfterMinReqVersion(String key) {
         Version version;
         if (MINIMAL_VERSION_NEURAL.containsKey(key)) {

src/main/java/org/opensearch/neuralsearch/rest/RestNeuralStatsAction.java

@@ -62,6 +62,21 @@ public class RestNeuralStatsAction extends BaseRestHandler {
      */
     public static final String INCLUDE_METADATA_PARAM = "include_metadata";
 
+    /**
+     * Query parameter name to include individual nodes data
+     */
+    public static final String INCLUDE_INDIVIDUAL_NODES_PARAM = "include_individual_nodes";
+
+    /**
+     * Query parameter name to include individual nodes data
+     */
+    public static final String INCLUDE_ALL_NODES_PARAM = "include_all_nodes";
+
+    /**
+     * Query parameter name to include individual nodes data
+     */
+    public static final String INCLUDE_INFO_PARAM = "include_info";
+
     /**
      * Regex for valid params, containing only alphanumeric, -, or _
      */
@@ -93,7 +108,15 @@ public class RestNeuralStatsAction extends BaseRestHandler {
         new Route(RestRequest.Method.GET, NEURAL_BASE_URI + "/stats/{stat}")
     );
 
-    private static final Set<String> RESPONSE_PARAMS = ImmutableSet.of(NODE_ID_PARAM, STAT_PARAM, INCLUDE_METADATA_PARAM, FLATTEN_PARAM);
+    private static final Set<String> RESPONSE_PARAMS = ImmutableSet.of(
+        NODE_ID_PARAM,
+        STAT_PARAM,
+        INCLUDE_METADATA_PARAM,
+        FLATTEN_PARAM,
+        INCLUDE_INDIVIDUAL_NODES_PARAM,
+        INCLUDE_ALL_NODES_PARAM,
+        INCLUDE_INFO_PARAM
+    );
 
     /**
      * Validates a param string if its under the max length and matches simple string pattern
@@ -172,6 +195,15 @@ private NeuralStatsInput createNeuralStatsInputFromRequestParams(RestRequest req
         boolean includeMetadata = request.paramAsBoolean(INCLUDE_METADATA_PARAM, false);
         neuralStatsInput.setIncludeMetadata(includeMetadata);
 
+        boolean includeIndividualNodes = request.paramAsBoolean(INCLUDE_INDIVIDUAL_NODES_PARAM, true);
+        neuralStatsInput.setIncludeIndividualNodes(includeIndividualNodes);
+
+        boolean includeAllNodes = request.paramAsBoolean(INCLUDE_ALL_NODES_PARAM, true);
+        neuralStatsInput.setIncludeAllNodes(includeAllNodes);
+
+        boolean includeInfo = request.paramAsBoolean(INCLUDE_INFO_PARAM, true);
+        neuralStatsInput.setIncludeInfo(includeInfo);
+
         // Process requested stats parameters
         processStatsRequestParameters(request, neuralStatsInput);
 
@@ -191,6 +223,8 @@ private void processStatsRequestParameters(RestRequest request, NeuralStatsInput
 
         String[] stats = optionalStats.get();
         Set<String> invalidStatNames = new HashSet<>();
+        boolean includeEvents = neuralStatsInput.isIncludeEvents();
+        boolean includeInfo = neuralStatsInput.isIncludeInfo();
 
         for (String stat : stats) {
             // Validate parameter
@@ -200,12 +234,12 @@ private void processStatsRequestParameters(RestRequest request, NeuralStatsInput
                 continue;
             }
 
-            if (InfoStatName.isValidName(normalizedStat)) {
+            if (includeInfo && InfoStatName.isValidName(normalizedStat)) {
                 InfoStatName infoStatName = InfoStatName.from(normalizedStat);
                 if (infoStatName.version().onOrBefore(minClusterVersion)) {
                     neuralStatsInput.getInfoStatNames().add(InfoStatName.from(normalizedStat));
                 }
-            } else if (EventStatName.isValidName(normalizedStat)) {
+            } else if (includeEvents && EventStatName.isValidName(normalizedStat)) {
                 EventStatName eventStatName = EventStatName.from(normalizedStat);
                 if (eventStatName.version().onOrBefore(minClusterVersion)) {
                     neuralStatsInput.getEventStatNames().add(EventStatName.from(normalizedStat));
@@ -224,24 +258,32 @@ private void processStatsRequestParameters(RestRequest request, NeuralStatsInput
 
     private void addAllStats(NeuralStatsInput neuralStatsInput, Version minVersion) {
         if (minVersion == Version.CURRENT) {
-            neuralStatsInput.getInfoStatNames().addAll(EnumSet.allOf(InfoStatName.class));
-            neuralStatsInput.getEventStatNames().addAll(EnumSet.allOf(EventStatName.class));
+            if (neuralStatsInput.isIncludeInfo()) {
+                neuralStatsInput.getInfoStatNames().addAll(EnumSet.allOf(InfoStatName.class));
+            }
+            if (neuralStatsInput.isIncludeEvents()) {
+                neuralStatsInput.getEventStatNames().addAll(EnumSet.allOf(EventStatName.class));
+            }
         } else {
             // Use a separate case here to save on version comparison if not necessary
-            neuralStatsInput.getInfoStatNames()
-                .addAll(
-                    EnumSet.allOf(InfoStatName.class)
-                        .stream()
-                        .filter(statName -> statName.version().onOrBefore(minVersion))
-                        .collect(Collectors.toCollection(() -> EnumSet.noneOf(InfoStatName.class)))
-                );
-            neuralStatsInput.getEventStatNames()
-                .addAll(
-                    EnumSet.allOf(EventStatName.class)
-                        .stream()
-                        .filter(statName -> statName.version().onOrBefore(minVersion))
-                        .collect(Collectors.toCollection(() -> EnumSet.noneOf(EventStatName.class)))
-                );
+            if (neuralStatsInput.isIncludeInfo()) {
+                neuralStatsInput.getInfoStatNames()
+                    .addAll(
+                        EnumSet.allOf(InfoStatName.class)
+                            .stream()
+                            .filter(statName -> statName.version().onOrBefore(minVersion))
+                            .collect(Collectors.toCollection(() -> EnumSet.noneOf(InfoStatName.class)))
+                    );
+            }
+            if (neuralStatsInput.isIncludeEvents()) {
+                neuralStatsInput.getEventStatNames()
+                    .addAll(
+                        EnumSet.allOf(EventStatName.class)
+                            .stream()
+                            .filter(statName -> statName.version().onOrBefore(minVersion))
+                            .collect(Collectors.toCollection(() -> EnumSet.noneOf(EventStatName.class)))
+                    );
+            }
         }
     }
 

src/main/java/org/opensearch/neuralsearch/stats/NeuralStatsInput.java

@@ -22,6 +22,8 @@
 import java.util.EnumSet;
 import java.util.List;
 
+import static org.opensearch.neuralsearch.common.MinClusterVersionUtil.isClusterOnOrAfterMinReqVersionForStatCategoryFiltering;
+
 /**
  * Entity class to hold input parameters for retrieving neural stats
  * Responsible for filtering statistics by node IDs, event statistic types, and info stat types.
@@ -60,6 +62,22 @@ public class NeuralStatsInput implements ToXContentObject, Writeable {
     @Setter
     private boolean flatten;
 
+    @Setter
+    /**
+     * Controls whether the response will include individual nodes
+     */
+    private boolean includeIndividualNodes;
+    @Setter
+    /**
+     * Controls whether the response will include aggregated nodes
+     */
+    private boolean includeAllNodes;
+    @Setter
+    /**
+     * Controls whether the response will include info nodes
+     */
+    private boolean includeInfo;
+
     /**
      * Builder constructor for creating NeuralStatsInput with specific filtering parameters.
      *
@@ -75,13 +93,19 @@ public NeuralStatsInput(
         EnumSet<EventStatName> eventStatNames,
         EnumSet<InfoStatName> infoStatNames,
         boolean includeMetadata,
-        boolean flatten
+        boolean flatten,
+        boolean includeIndividualNodes,
+        boolean includeAllNodes,
+        boolean includeInfo
     ) {
         this.nodeIds = nodeIds;
         this.eventStatNames = eventStatNames;
         this.infoStatNames = infoStatNames;
         this.includeMetadata = includeMetadata;
         this.flatten = flatten;
+        this.includeIndividualNodes = includeIndividualNodes;
+        this.includeAllNodes = includeAllNodes;
+        this.includeInfo = includeInfo;
     }
 
     /**
@@ -94,6 +118,9 @@ public NeuralStatsInput() {
         this.infoStatNames = EnumSet.noneOf(InfoStatName.class);
         this.includeMetadata = false;
         this.flatten = false;
+        this.includeIndividualNodes = true;
+        this.includeAllNodes = true;
+        this.includeInfo = true;
     }
 
     /**
@@ -108,6 +135,15 @@ public NeuralStatsInput(StreamInput input) throws IOException {
         infoStatNames = input.readOptionalEnumSet(InfoStatName.class);
         includeMetadata = input.readBoolean();
         flatten = input.readBoolean();
+        if (isClusterOnOrAfterMinReqVersionForStatCategoryFiltering()) {
+            includeIndividualNodes = input.readBoolean();
+            includeAllNodes = input.readBoolean();
+            includeInfo = input.readBoolean();
+        } else {
+            includeIndividualNodes = true;
+            includeAllNodes = true;
+            includeInfo = true;
+        }
     }
 
     /**
@@ -123,6 +159,11 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalEnumSet(infoStatNames);
         out.writeBoolean(includeMetadata);
         out.writeBoolean(flatten);
+        if (isClusterOnOrAfterMinReqVersionForStatCategoryFiltering()) {
+            out.writeBoolean(includeIndividualNodes);
+            out.writeBoolean(includeAllNodes);
+            out.writeBoolean(includeInfo);
+        }
     }
 
     /**
@@ -147,7 +188,19 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
         }
         builder.field(RestNeuralStatsAction.INCLUDE_METADATA_PARAM, includeMetadata);
         builder.field(RestNeuralStatsAction.FLATTEN_PARAM, flatten);
+        builder.field(RestNeuralStatsAction.INCLUDE_INDIVIDUAL_NODES_PARAM, includeIndividualNodes);
+        builder.field(RestNeuralStatsAction.INCLUDE_ALL_NODES_PARAM, includeAllNodes);
+        builder.field(RestNeuralStatsAction.INCLUDE_INFO_PARAM, includeInfo);
         builder.endObject();
         return builder;
     }
+
+    /**
+     * Helper to determine if we should fetch event stats or if we can skip them
+     * If we exclude both individual and all nodes, then there is no need to fetch any specific stats from nodes
+     * @return whether we need to fetch event stats
+     */
+    public boolean isIncludeEvents() {
+        return this.isIncludeAllNodes() || this.isIncludeIndividualNodes();
+    }
 }