Skip to content

Commit 8f4ae63

Browse files
benwtrentbenwtrent
authored
[ML][Data Frame] add node attr to GET _stats (#43842)
* [ML][Data Frame] add node attr to GET _stats * addressing testing issues with node.attributes
1 parent d1c6fb8 commit 8f4ae63

File tree

14 files changed

+631
-24
lines changed

14 files changed

+631
-24
lines changed

‎client/rest-high-level/src/main/java/org/elasticsearch/client/dataframe/transforms/DataFrameTransformState.java

+15-3
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ public class DataFrameTransformState {
4343
private static final ParseField CHECKPOINT = new ParseField("checkpoint");
4444
private static final ParseField REASON = new ParseField("reason");
4545
private static final ParseField PROGRESS = new ParseField("progress");
46+
private static final ParseField NODE = new ParseField("node");
4647

4748
@SuppressWarnings("unchecked")
4849
public static final ConstructingObjectParser<DataFrameTransformState, Void> PARSER =
@@ -52,7 +53,8 @@ public class DataFrameTransformState {
5253
(Map<String, Object>) args[2],
5354
(long) args[3],
5455
(String) args[4],
55-
(DataFrameTransformProgress) args[5]));
56+
(DataFrameTransformProgress) args[5],
57+
(NodeAttributes) args[6]));
5658

5759
static {
5860
PARSER.declareField(constructorArg(), p -> DataFrameTransformTaskState.fromString(p.text()), TASK_STATE, ValueType.STRING);
@@ -61,6 +63,7 @@ public class DataFrameTransformState {
6163
PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), CHECKPOINT);
6264
PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), REASON);
6365
PARSER.declareField(optionalConstructorArg(), DataFrameTransformProgress::fromXContent, PROGRESS, ValueType.OBJECT);
66+
PARSER.declareField(optionalConstructorArg(), NodeAttributes.PARSER::apply, NODE, ValueType.OBJECT);
6467
}
6568

6669
public static DataFrameTransformState fromXContent(XContentParser parser) throws IOException {
@@ -73,19 +76,22 @@ public static DataFrameTransformState fromXContent(XContentParser parser) throws
7376
private final Map<String, Object> currentPosition;
7477
private final String reason;
7578
private final DataFrameTransformProgress progress;
79+
private final NodeAttributes node;
7680

7781
public DataFrameTransformState(DataFrameTransformTaskState taskState,
7882
IndexerState indexerState,
7983
@Nullable Map<String, Object> position,
8084
long checkpoint,
8185
@Nullable String reason,
82-
@Nullable DataFrameTransformProgress progress) {
86+
@Nullable DataFrameTransformProgress progress,
87+
@Nullable NodeAttributes node) {
8388
this.taskState = taskState;
8489
this.indexerState = indexerState;
8590
this.currentPosition = position == null ? null : Collections.unmodifiableMap(new LinkedHashMap<>(position));
8691
this.checkpoint = checkpoint;
8792
this.reason = reason;
8893
this.progress = progress;
94+
this.node = node;
8995
}
9096

9197
public IndexerState getIndexerState() {
@@ -115,6 +121,11 @@ public DataFrameTransformProgress getProgress() {
115121
return progress;
116122
}
117123

124+
@Nullable
125+
public NodeAttributes getNode() {
126+
return node;
127+
}
128+
118129
@Override
119130
public boolean equals(Object other) {
120131
if (this == other) {
@@ -132,12 +143,13 @@ public boolean equals(Object other) {
132143
Objects.equals(this.currentPosition, that.currentPosition) &&
133144
Objects.equals(this.progress, that.progress) &&
134145
this.checkpoint == that.checkpoint &&
146+
Objects.equals(this.node, that.node) &&
135147
Objects.equals(this.reason, that.reason);
136148
}
137149

138150
@Override
139151
public int hashCode() {
140-
return Objects.hash(taskState, indexerState, currentPosition, checkpoint, reason, progress);
152+
return Objects.hash(taskState, indexerState, currentPosition, checkpoint, reason, progress, node);
141153
}
142154

143155
}

‎client/rest-high-level/src/main/java/org/elasticsearch/client/dataframe/transforms/NodeAttributes.java

+156
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.client.dataframe.transforms;
20+
21+
import org.elasticsearch.common.ParseField;
22+
import org.elasticsearch.common.Strings;
23+
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
24+
import org.elasticsearch.common.xcontent.ObjectParser;
25+
import org.elasticsearch.common.xcontent.ToXContentObject;
26+
import org.elasticsearch.common.xcontent.XContentBuilder;
27+
28+
import java.io.IOException;
29+
import java.util.Collections;
30+
import java.util.Map;
31+
import java.util.Objects;
32+
33+
/**
34+
* A Pojo class containing an Elastic Node's attributes
35+
*/
36+
public class NodeAttributes implements ToXContentObject {
37+
38+
public static final ParseField ID = new ParseField("id");
39+
public static final ParseField NAME = new ParseField("name");
40+
public static final ParseField EPHEMERAL_ID = new ParseField("ephemeral_id");
41+
public static final ParseField TRANSPORT_ADDRESS = new ParseField("transport_address");
42+
public static final ParseField ATTRIBUTES = new ParseField("attributes");
43+
44+
@SuppressWarnings("unchecked")
45+
public static final ConstructingObjectParser<NodeAttributes, Void> PARSER =
46+
new ConstructingObjectParser<>("node", true,
47+
(a) -> {
48+
int i = 0;
49+
String id = (String) a[i++];
50+
String name = (String) a[i++];
51+
String ephemeralId = (String) a[i++];
52+
String transportAddress = (String) a[i++];
53+
Map<String, String> attributes = (Map<String, String>) a[i];
54+
return new NodeAttributes(id, name, ephemeralId, transportAddress, attributes);
55+
});
56+
57+
static {
58+
PARSER.declareString(ConstructingObjectParser.constructorArg(), ID);
59+
PARSER.declareString(ConstructingObjectParser.constructorArg(), NAME);
60+
PARSER.declareString(ConstructingObjectParser.constructorArg(), EPHEMERAL_ID);
61+
PARSER.declareString(ConstructingObjectParser.constructorArg(), TRANSPORT_ADDRESS);
62+
PARSER.declareField(ConstructingObjectParser.constructorArg(),
63+
(p, c) -> p.mapStrings(),
64+
ATTRIBUTES,
65+
ObjectParser.ValueType.OBJECT);
66+
}
67+
68+
private final String id;
69+
private final String name;
70+
private final String ephemeralId;
71+
private final String transportAddress;
72+
private final Map<String, String> attributes;
73+
74+
public NodeAttributes(String id, String name, String ephemeralId, String transportAddress, Map<String, String> attributes) {
75+
this.id = id;
76+
this.name = name;
77+
this.ephemeralId = ephemeralId;
78+
this.transportAddress = transportAddress;
79+
this.attributes = Collections.unmodifiableMap(attributes);
80+
}
81+
82+
/**
83+
* The unique identifier of the node.
84+
*/
85+
public String getId() {
86+
return id;
87+
}
88+
89+
/**
90+
* The node name.
91+
*/
92+
public String getName() {
93+
return name;
94+
}
95+
96+
/**
97+
* The ephemeral id of the node.
98+
*/
99+
public String getEphemeralId() {
100+
return ephemeralId;
101+
}
102+
103+
/**
104+
* The host and port where transport HTTP connections are accepted.
105+
*/
106+
public String getTransportAddress() {
107+
return transportAddress;
108+
}
109+
110+
/**
111+
* Additional attributes related to this node
112+
*/
113+
public Map<String, String> getAttributes() {
114+
return attributes;
115+
}
116+
117+
@Override
118+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
119+
builder.startObject();
120+
builder.field(ID.getPreferredName(), id);
121+
builder.field(NAME.getPreferredName(), name);
122+
builder.field(EPHEMERAL_ID.getPreferredName(), ephemeralId);
123+
builder.field(TRANSPORT_ADDRESS.getPreferredName(), transportAddress);
124+
builder.field(ATTRIBUTES.getPreferredName(), attributes);
125+
builder.endObject();
126+
return builder;
127+
}
128+
129+
@Override
130+
public int hashCode() {
131+
return Objects.hash(id, name, ephemeralId, transportAddress, attributes);
132+
}
133+
134+
@Override
135+
public boolean equals(Object other) {
136+
if (this == other) {
137+
return true;
138+
}
139+
140+
if (other == null || getClass() != other.getClass()) {
141+
return false;
142+
}
143+
144+
NodeAttributes that = (NodeAttributes) other;
145+
return Objects.equals(id, that.id) &&
146+
Objects.equals(name, that.name) &&
147+
Objects.equals(ephemeralId, that.ephemeralId) &&
148+
Objects.equals(transportAddress, that.transportAddress) &&
149+
Objects.equals(attributes, that.attributes);
150+
}
151+
152+
@Override
153+
public String toString() {
154+
return Strings.toString(this);
155+
}
156+
}

‎client/rest-high-level/src/test/java/org/elasticsearch/client/dataframe/transforms/DataFrameTransformStateTests.java

+9-2
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
package org.elasticsearch.client.dataframe.transforms;
2121

2222
import org.elasticsearch.client.core.IndexerState;
23+
import org.elasticsearch.common.xcontent.ToXContent;
2324
import org.elasticsearch.common.xcontent.XContentBuilder;
2425
import org.elasticsearch.test.ESTestCase;
2526

@@ -37,7 +38,8 @@ public void testFromXContent() throws IOException {
3738
DataFrameTransformStateTests::toXContent,
3839
DataFrameTransformState::fromXContent)
3940
.supportsUnknownFields(true)
40-
.randomFieldsExcludeFilter(field -> field.equals("current_position"))
41+
.randomFieldsExcludeFilter(field -> field.equals("current_position") ||
42+
field.equals("node.attributes"))
4143
.test();
4244
}
4345

@@ -47,7 +49,8 @@ public static DataFrameTransformState randomDataFrameTransformState() {
4749
randomPositionMap(),
4850
randomLongBetween(0,10),
4951
randomBoolean() ? null : randomAlphaOfLength(10),
50-
randomBoolean() ? null : DataFrameTransformProgressTests.randomInstance());
52+
randomBoolean() ? null : DataFrameTransformProgressTests.randomInstance(),
53+
randomBoolean() ? null : NodeAttributesTests.createRandom());
5154
}
5255

5356
public static void toXContent(DataFrameTransformState state, XContentBuilder builder) throws IOException {
@@ -65,6 +68,10 @@ public static void toXContent(DataFrameTransformState state, XContentBuilder bui
6568
builder.field("progress");
6669
DataFrameTransformProgressTests.toXContent(state.getProgress(), builder);
6770
}
71+
if (state.getNode() != null) {
72+
builder.field("node");
73+
state.getNode().toXContent(builder, ToXContent.EMPTY_PARAMS);
74+
}
6875
builder.endObject();
6976
}
7077

‎client/rest-high-level/src/test/java/org/elasticsearch/client/dataframe/transforms/NodeAttributesTests.java

+64
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.client.dataframe.transforms;
20+
21+
import org.elasticsearch.common.xcontent.XContentParser;
22+
import org.elasticsearch.test.AbstractXContentTestCase;
23+
24+
import java.io.IOException;
25+
import java.util.HashMap;
26+
import java.util.Map;
27+
import java.util.function.Predicate;
28+
29+
public class NodeAttributesTests extends AbstractXContentTestCase<NodeAttributes> {
30+
31+
public static NodeAttributes createRandom() {
32+
int numberOfAttributes = randomIntBetween(1, 10);
33+
Map<String, String> attributes = new HashMap<>(numberOfAttributes);
34+
for(int i = 0; i < numberOfAttributes; i++) {
35+
String val = randomAlphaOfLength(10);
36+
attributes.put("key-"+i, val);
37+
}
38+
return new NodeAttributes(randomAlphaOfLength(10),
39+
randomAlphaOfLength(10),
40+
randomAlphaOfLength(10),
41+
randomAlphaOfLength(10),
42+
attributes);
43+
}
44+
45+
@Override
46+
protected NodeAttributes createTestInstance() {
47+
return createRandom();
48+
}
49+
50+
@Override
51+
protected NodeAttributes doParseInstance(XContentParser parser) throws IOException {
52+
return NodeAttributes.PARSER.parse(parser, null);
53+
}
54+
55+
@Override
56+
protected Predicate<String> getRandomFieldsExcludeFilter() {
57+
return field -> !field.isEmpty();
58+
}
59+
60+
@Override
61+
protected boolean supportsUnknownFields() {
62+
return true;
63+
}
64+
}

‎client/rest-high-level/src/test/java/org/elasticsearch/client/dataframe/transforms/hlrc/DataFrameTransformStateAndStatsTests.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ protected boolean supportsUnknownFields() {
6464

6565
@Override
6666
protected Predicate<String> getRandomFieldsExcludeFilter() {
67-
return field -> field.equals("state.current_position");
67+
return field -> field.equals("state.current_position") || field.equals("state.node") || field.equals("state.node.attributes");
6868
}
6969
}
7070

0 commit comments

Comments
 (0)