Skip to content

Commit 0426b26

Browse files
q-andyq-andy
committed
Add include_category flags
Signed-off-by: Andy Qin <[email protected]>
1 parent 575f774 commit 0426b26

File tree

9 files changed

+254
-41
lines changed

9 files changed

+254
-41
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1717
- [Stats] Add stats for text chunking processor algorithms ([#1308](https://github.com/opensearch-project/neural-search/pull/1308))
1818
- Support custom weights in RRF normalization processor ([#1322](https://github.com/opensearch-project/neural-search/pull/1322))
1919
- [Stats] Add stats tracking for semantic highlighting ([#1327](https://github.com/opensearch-project/neural-search/pull/1327))
20-
20+
- [Stats] Add `include_individual_nodes` parameter to stats API ([#1360](https://github.com/opensearch-project/neural-search/pull/1360))
2121
### Bug Fixes
2222
- 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))
2323
- 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))

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

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,21 @@ public class RestNeuralStatsAction extends BaseRestHandler {
6060
*/
6161
public static final String INCLUDE_METADATA_PARAM = "include_metadata";
6262

63+
/**
64+
* Query parameter name to include individual nodes data
65+
*/
66+
public static final String INCLUDE_INDIVIDUAL_NODES_PARAM = "include_individual_nodes";
67+
68+
/**
69+
* Query parameter name to include individual nodes data
70+
*/
71+
public static final String INCLUDE_ALL_NODES_PARAM = "include_all_nodes";
72+
73+
/**
74+
* Query parameter name to include individual nodes data
75+
*/
76+
public static final String INCLUDE_INFO = "include_info";
77+
6378
/**
6479
* Regex for valid params, containing only alphanumeric, -, or _
6580
*/
@@ -91,7 +106,13 @@ public class RestNeuralStatsAction extends BaseRestHandler {
91106
new Route(RestRequest.Method.GET, NEURAL_BASE_URI + "/stats/{stat}")
92107
);
93108

94-
private static final Set<String> RESPONSE_PARAMS = ImmutableSet.of(NODE_ID_PARAM, STAT_PARAM, INCLUDE_METADATA_PARAM, FLATTEN_PARAM);
109+
private static final Set<String> RESPONSE_PARAMS = ImmutableSet.of(
110+
NODE_ID_PARAM,
111+
STAT_PARAM,
112+
INCLUDE_METADATA_PARAM,
113+
FLATTEN_PARAM,
114+
INCLUDE_INDIVIDUAL_NODES_PARAM
115+
);
95116

96117
/**
97118
* Validates a param string if its under the max length and matches simple string pattern
@@ -169,6 +190,9 @@ private NeuralStatsInput createNeuralStatsInputFromRequestParams(RestRequest req
169190
boolean includeMetadata = request.paramAsBoolean(INCLUDE_METADATA_PARAM, false);
170191
neuralStatsInput.setIncludeMetadata(includeMetadata);
171192

193+
boolean includeIndividualNodes = request.paramAsBoolean(INCLUDE_INDIVIDUAL_NODES_PARAM, true);
194+
neuralStatsInput.setIncludeIndividualNodes(includeIndividualNodes);
195+
172196
// Process requested stats parameters
173197
processStatsRequestParameters(request, neuralStatsInput);
174198

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

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,22 @@ public class NeuralStatsInput implements ToXContentObject, Writeable {
6060
@Setter
6161
private boolean flatten;
6262

63+
@Setter
64+
/**
65+
* Controls whether the response will include individual nodes
66+
*/
67+
private boolean includeIndividualNodes;
68+
@Setter
69+
/**
70+
* Controls whether the response will include aggregated nodes
71+
*/
72+
private boolean includeAllNodes;
73+
@Setter
74+
/**
75+
* Controls whether the response will include info nodes
76+
*/
77+
private boolean includeInfo;
78+
6379
/**
6480
* Builder constructor for creating NeuralStatsInput with specific filtering parameters.
6581
*
@@ -75,13 +91,19 @@ public NeuralStatsInput(
7591
EnumSet<EventStatName> eventStatNames,
7692
EnumSet<InfoStatName> infoStatNames,
7793
boolean includeMetadata,
78-
boolean flatten
94+
boolean flatten,
95+
boolean includeIndividualNodes,
96+
boolean includeAllNodes,
97+
boolean includeInfo
7998
) {
8099
this.nodeIds = nodeIds;
81100
this.eventStatNames = eventStatNames;
82101
this.infoStatNames = infoStatNames;
83102
this.includeMetadata = includeMetadata;
84103
this.flatten = flatten;
104+
this.includeIndividualNodes = includeIndividualNodes;
105+
this.includeAllNodes = includeAllNodes;
106+
this.includeInfo = includeInfo;
85107
}
86108

87109
/**
@@ -94,6 +116,9 @@ public NeuralStatsInput() {
94116
this.infoStatNames = EnumSet.noneOf(InfoStatName.class);
95117
this.includeMetadata = false;
96118
this.flatten = false;
119+
this.includeIndividualNodes = true;
120+
this.includeAllNodes = true;
121+
this.includeInfo = true;
97122
}
98123

99124
/**
@@ -108,6 +133,9 @@ public NeuralStatsInput(StreamInput input) throws IOException {
108133
infoStatNames = input.readOptionalEnumSet(InfoStatName.class);
109134
includeMetadata = input.readBoolean();
110135
flatten = input.readBoolean();
136+
includeIndividualNodes = input.readBoolean();
137+
includeAllNodes = input.readBoolean();
138+
includeInfo = input.readBoolean();
111139
}
112140

113141
/**
@@ -123,6 +151,9 @@ public void writeTo(StreamOutput out) throws IOException {
123151
out.writeOptionalEnumSet(infoStatNames);
124152
out.writeBoolean(includeMetadata);
125153
out.writeBoolean(flatten);
154+
out.writeBoolean(includeIndividualNodes);
155+
out.writeBoolean(includeAllNodes);
156+
out.writeBoolean(includeInfo);
126157
}
127158

128159
/**
@@ -147,6 +178,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
147178
}
148179
builder.field(RestNeuralStatsAction.INCLUDE_METADATA_PARAM, includeMetadata);
149180
builder.field(RestNeuralStatsAction.FLATTEN_PARAM, flatten);
181+
builder.field(RestNeuralStatsAction.INCLUDE_INDIVIDUAL_NODES_PARAM, includeIndividualNodes);
182+
builder.field(RestNeuralStatsAction.INCLUDE_ALL_NODES_PARAM, includeAllNodes);
183+
builder.field(RestNeuralStatsAction.INCLUDE_INFO, includeInfo);
150184
builder.endObject();
151185
return builder;
152186
}

‎src/main/java/org/opensearch/neuralsearch/transport/NeuralStatsResponse.java

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ public class NeuralStatsResponse extends BaseNodesResponse<NeuralStatsNodeRespon
3434
private Map<String, Map<String, StatSnapshot<?>>> nodeIdToNodeEventStats;
3535
private boolean flatten;
3636
private boolean includeMetadata;
37+
private boolean includeIndividualNodes;
38+
private boolean includeAllNodes;
39+
private boolean includeInfo;
3740

3841
/**
3942
* Constructor
@@ -53,6 +56,9 @@ public NeuralStatsResponse(StreamInput in) throws IOException {
5356
this.nodeIdToNodeEventStats = castedNodeIdToNodeEventStats;
5457
this.flatten = in.readBoolean();
5558
this.includeMetadata = in.readBoolean();
59+
this.includeIndividualNodes = in.readBoolean();
60+
this.includeAllNodes = in.readBoolean();
61+
this.includeInfo = in.readBoolean();
5662
}
5763

5864
/**
@@ -75,14 +81,20 @@ public NeuralStatsResponse(
7581
Map<String, StatSnapshot<?>> aggregatedNodeStats,
7682
Map<String, Map<String, StatSnapshot<?>>> nodeIdToNodeEventStats,
7783
boolean flatten,
78-
boolean includeMetadata
84+
boolean includeMetadata,
85+
boolean includeIndividualNodes,
86+
boolean includeAllNodes,
87+
boolean includeInfo
7988
) {
8089
super(clusterName, nodes, failures);
8190
this.infoStats = infoStats;
8291
this.aggregatedNodeStats = aggregatedNodeStats;
8392
this.nodeIdToNodeEventStats = nodeIdToNodeEventStats;
8493
this.flatten = flatten;
8594
this.includeMetadata = includeMetadata;
95+
this.includeIndividualNodes = includeIndividualNodes;
96+
this.includeAllNodes = includeAllNodes;
97+
this.includeInfo = includeInfo;
8698
}
8799

88100
@Override
@@ -97,6 +109,9 @@ public void writeTo(StreamOutput out) throws IOException {
97109
out.writeMap(downcastedNodeIdToNodeEventStats);
98110
out.writeBoolean(flatten);
99111
out.writeBoolean(includeMetadata);
112+
out.writeBoolean(includeIndividualNodes);
113+
out.writeBoolean(includeAllNodes);
114+
out.writeBoolean(includeInfo);
100115
}
101116

102117
@Override
@@ -111,20 +126,26 @@ public List<NeuralStatsNodeResponse> readNodesFrom(StreamInput in) throws IOExce
111126

112127
@Override
113128
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
114-
Map<String, Object> formattedInfoStats = formatStats(infoStats);
115-
builder.startObject(INFO_KEY_PREFIX);
116-
builder.mapContents(formattedInfoStats);
117-
builder.endObject();
118-
119-
Map<String, Object> formattedAggregatedNodeStats = formatStats(aggregatedNodeStats);
120-
builder.startObject(AGGREGATED_NODES_KEY_PREFIX);
121-
builder.mapContents(formattedAggregatedNodeStats);
122-
builder.endObject();
123-
124-
Map<String, Object> formattedNodeEventStats = formatNodeEventStats(nodeIdToNodeEventStats);
125-
builder.startObject(NODES_KEY_PREFIX);
126-
builder.mapContents(formattedNodeEventStats);
127-
builder.endObject();
129+
if (includeInfo) {
130+
Map<String, Object> formattedInfoStats = formatStats(infoStats);
131+
builder.startObject(INFO_KEY_PREFIX);
132+
builder.mapContents(formattedInfoStats);
133+
builder.endObject();
134+
}
135+
136+
if (includeAllNodes) {
137+
Map<String, Object> formattedAggregatedNodeStats = formatStats(aggregatedNodeStats);
138+
builder.startObject(AGGREGATED_NODES_KEY_PREFIX);
139+
builder.mapContents(formattedAggregatedNodeStats);
140+
builder.endObject();
141+
}
142+
143+
if (includeIndividualNodes) {
144+
Map<String, Object> formattedNodeEventStats = formatNodeEventStats(nodeIdToNodeEventStats);
145+
builder.startObject(NODES_KEY_PREFIX);
146+
builder.mapContents(formattedNodeEventStats);
147+
builder.endObject();
148+
}
128149

129150
return builder;
130151
}

‎src/main/java/org/opensearch/neuralsearch/transport/NeuralStatsTransportAction.java

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -77,9 +77,6 @@ protected NeuralStatsResponse newResponse(
7777
List<NeuralStatsNodeResponse> responses,
7878
List<FailedNodeException> failures
7979
) {
80-
// Final object that will hold the stats in format Map<ResponsePath, Value>
81-
Map<String, StatSnapshot<?>> resultStats = new HashMap<>();
82-
8380
// Convert node level stats to map
8481
Map<String, Map<String, StatSnapshot<?>>> nodeIdToEventStats = processorNodeEventStatsIntoMap(responses);
8582

@@ -89,13 +86,16 @@ protected NeuralStatsResponse newResponse(
8986
request.getNeuralStatsInput().getEventStatNames()
9087
);
9188

92-
// Get info stats
93-
Map<InfoStatName, StatSnapshot<?>> infoStats = infoStatsManager.getStats(request.getNeuralStatsInput().getInfoStatNames());
89+
Map<String, StatSnapshot<?>> flatInfoStats = new HashMap<>();
90+
if (request.getNeuralStatsInput().isIncludeInfo()) {
91+
// Get info stats
92+
Map<InfoStatName, StatSnapshot<?>> infoStats = infoStatsManager.getStats(request.getNeuralStatsInput().getInfoStatNames());
9493

95-
// Convert stat name keys into flat path strings
96-
Map<String, StatSnapshot<?>> flatInfoStats = infoStats.entrySet()
97-
.stream()
98-
.collect(Collectors.toMap(entry -> entry.getKey().getFullPath(), Map.Entry::getValue));
94+
// Convert stat name keys into flat path strings
95+
flatInfoStats = infoStats.entrySet()
96+
.stream()
97+
.collect(Collectors.toMap(entry -> entry.getKey().getFullPath(), Map.Entry::getValue));
98+
}
9999

100100
return new NeuralStatsResponse(
101101
clusterService.getClusterName(),
@@ -105,7 +105,10 @@ protected NeuralStatsResponse newResponse(
105105
aggregatedNodeStats,
106106
nodeIdToEventStats,
107107
request.getNeuralStatsInput().isFlatten(),
108-
request.getNeuralStatsInput().isIncludeMetadata()
108+
request.getNeuralStatsInput().isIncludeMetadata(),
109+
request.getNeuralStatsInput().isIncludeIndividualNodes(),
110+
request.getNeuralStatsInput().isIncludeAllNodes(),
111+
request.getNeuralStatsInput().isIncludeInfo()
109112
);
110113
}
111114

‎src/test/java/org/opensearch/neuralsearch/rest/RestNeuralStatsActionIT.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.opensearch.client.Request;
1111
import org.opensearch.client.ResponseException;
1212
import org.opensearch.core.rest.RestStatus;
13+
import org.opensearch.core.xcontent.MediaTypeRegistry;
1314
import org.opensearch.neuralsearch.BaseNeuralSearchIT;
1415
import org.opensearch.neuralsearch.plugin.NeuralSearch;
1516
import org.opensearch.neuralsearch.settings.NeuralSearchSettings;
@@ -86,6 +87,7 @@ public void test_textEmbedding() throws Exception {
8687
String responseBody;
8788
Map<String, Object> stats;
8889
List<Map<String, Object>> nodesStats;
90+
Map<String, Object> aggregatedNodesStats;
8991

9092
responseBody = executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>());
9193
stats = parseInfoStatsResponse(responseBody);
@@ -95,6 +97,27 @@ public void test_textEmbedding() throws Exception {
9597
assertEquals(3, getNestedValue(nodesStats.getFirst(), EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS));
9698
assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS));
9799

100+
// Aggregated stats should still exist even if include_individual_nodes is false
101+
Map<String, String> params = new HashMap<>();
102+
params.put(RestNeuralStatsAction.INCLUDE_INDIVIDUAL_NODES_PARAM, "false");
103+
104+
// Call stats again with custom params
105+
responseBody = executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>(), params);
106+
stats = parseInfoStatsResponse(responseBody);
107+
aggregatedNodesStats = parseAggregatedNodeStatsResponse(responseBody);
108+
109+
// Individual nodes shouldn't exist
110+
@SuppressWarnings("unchecked")
111+
Map<String, Object> individualNodesMap = (Map<String, Object>) createParser(
112+
MediaTypeRegistry.getDefaultMediaType().xContent(),
113+
responseBody
114+
).map().get("nodes");
115+
assertNull(individualNodesMap);
116+
117+
// Parse aggregated nodes stats json
118+
assertEquals(3, getNestedValue(aggregatedNodesStats, EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS));
119+
assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS));
120+
98121
// Reset stats
99122
updateClusterSettings("plugins.neural_search.stats_enabled", false);
100123
updateClusterSettings("plugins.neural_search.stats_enabled", true);

‎src/test/java/org/opensearch/neuralsearch/rest/RestNeuralStatsActionTests.java

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public void tearDown() throws Exception {
7171
client.close();
7272
}
7373

74-
public void test_execute() throws Exception {
74+
public void test_execute_defaultParams() throws Exception {
7575
when(settingsAccessor.isStatsEnabled()).thenReturn(true);
7676
RestNeuralStatsAction restNeuralStatsAction = new RestNeuralStatsAction(settingsAccessor);
7777

@@ -84,6 +84,32 @@ public void test_execute() throws Exception {
8484
NeuralStatsInput capturedInput = argumentCaptor.getValue().getNeuralStatsInput();
8585
assertEquals(capturedInput.getEventStatNames(), EnumSet.allOf(EventStatName.class));
8686
assertEquals(capturedInput.getInfoStatNames(), EnumSet.allOf(InfoStatName.class));
87+
assertFalse(capturedInput.isFlatten());
88+
assertFalse(capturedInput.isIncludeMetadata());
89+
assertTrue(capturedInput.isIncludeIndividualNodes());
90+
}
91+
92+
public void test_execute_customParams() throws Exception {
93+
when(settingsAccessor.isStatsEnabled()).thenReturn(true);
94+
RestNeuralStatsAction restNeuralStatsAction = new RestNeuralStatsAction(settingsAccessor);
95+
96+
Map<String, String> params = new HashMap<>();
97+
params.put(RestNeuralStatsAction.FLATTEN_PARAM, "true");
98+
params.put(RestNeuralStatsAction.INCLUDE_METADATA_PARAM, "true");
99+
params.put(RestNeuralStatsAction.INCLUDE_INDIVIDUAL_NODES_PARAM, "false");
100+
RestRequest request = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY).withParams(params).build();
101+
102+
restNeuralStatsAction.handleRequest(request, channel, client);
103+
104+
ArgumentCaptor<NeuralStatsRequest> argumentCaptor = ArgumentCaptor.forClass(NeuralStatsRequest.class);
105+
verify(client, times(1)).execute(eq(NeuralStatsAction.INSTANCE), argumentCaptor.capture(), any());
106+
107+
NeuralStatsInput capturedInput = argumentCaptor.getValue().getNeuralStatsInput();
108+
assertEquals(capturedInput.getEventStatNames(), EnumSet.allOf(EventStatName.class));
109+
assertEquals(capturedInput.getInfoStatNames(), EnumSet.allOf(InfoStatName.class));
110+
assertTrue(capturedInput.isFlatten());
111+
assertTrue(capturedInput.isIncludeMetadata());
112+
assertFalse(capturedInput.isIncludeIndividualNodes());
87113
}
88114

89115
public void test_handleRequest_disabledForbidden() throws Exception {

0 commit comments

Comments
 (0)