Skip to content

Commit 9bf8b5a

Browse files
benwtrentbenwtrent
authored
[ML] Adds progress reporting for transforms (#41278)
* [ML] Adds progress reporting for transforms * fixing after master merge * Addressing PR comments * removing unused imports * Adjusting afterKey handling and percentage to be 100* * Making sure it is a linked hashmap for serialization * removing unused import * addressing PR comments * removing unused import * simplifying code, only storing total docs and decrementing * adjusting for rewrite * removing initial progress gathering from executor
1 parent f26addc commit 9bf8b5a

File tree

28 files changed

+1284
-264
lines changed

28 files changed

+1284
-264
lines changed

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

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
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+
20+
package org.elasticsearch.client.dataframe.transforms;
21+
22+
import org.elasticsearch.common.ParseField;
23+
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
24+
import org.elasticsearch.common.xcontent.XContentParser;
25+
26+
import java.util.Objects;
27+
28+
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
29+
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
30+
31+
public class DataFrameTransformProgress {
32+
33+
public static final ParseField TOTAL_DOCS = new ParseField("total_docs");
34+
public static final ParseField DOCS_REMAINING = new ParseField("docs_remaining");
35+
public static final ParseField PERCENT_COMPLETE = new ParseField("percent_complete");
36+
37+
public static final ConstructingObjectParser<DataFrameTransformProgress, Void> PARSER = new ConstructingObjectParser<>(
38+
"data_frame_transform_progress",
39+
true,
40+
a -> new DataFrameTransformProgress((Long) a[0], (Long)a[1], (Double)a[2]));
41+
42+
static {
43+
PARSER.declareLong(constructorArg(), TOTAL_DOCS);
44+
PARSER.declareLong(optionalConstructorArg(), DOCS_REMAINING);
45+
PARSER.declareDouble(optionalConstructorArg(), PERCENT_COMPLETE);
46+
}
47+
48+
public static DataFrameTransformProgress fromXContent(XContentParser parser) {
49+
return PARSER.apply(parser, null);
50+
}
51+
52+
private final long totalDocs;
53+
private final long remainingDocs;
54+
private final double percentComplete;
55+
56+
public DataFrameTransformProgress(long totalDocs, Long remainingDocs, double percentComplete) {
57+
this.totalDocs = totalDocs;
58+
this.remainingDocs = remainingDocs == null ? totalDocs : remainingDocs;
59+
this.percentComplete = percentComplete;
60+
}
61+
62+
public double getPercentComplete() {
63+
return percentComplete;
64+
}
65+
66+
public long getTotalDocs() {
67+
return totalDocs;
68+
}
69+
70+
public long getRemainingDocs() {
71+
return remainingDocs;
72+
}
73+
74+
@Override
75+
public boolean equals(Object other) {
76+
if (other == this) {
77+
return true;
78+
}
79+
80+
if (other == null || other.getClass() != getClass()) {
81+
return false;
82+
}
83+
84+
DataFrameTransformProgress that = (DataFrameTransformProgress) other;
85+
return Objects.equals(this.remainingDocs, that.remainingDocs)
86+
&& Objects.equals(this.totalDocs, that.totalDocs)
87+
&& Objects.equals(this.percentComplete, that.percentComplete);
88+
}
89+
90+
@Override
91+
public int hashCode(){
92+
return Objects.hash(remainingDocs, totalDocs, percentComplete);
93+
}
94+
}

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

Lines changed: 23 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -23,16 +23,14 @@
2323
import org.elasticsearch.common.Nullable;
2424
import org.elasticsearch.common.ParseField;
2525
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
26-
import org.elasticsearch.common.xcontent.ObjectParser;
26+
import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
2727
import org.elasticsearch.common.xcontent.XContentParser;
2828

2929
import java.io.IOException;
3030
import java.util.Collections;
31-
import java.util.HashMap;
31+
import java.util.LinkedHashMap;
3232
import java.util.Map;
3333
import java.util.Objects;
34-
import java.util.SortedMap;
35-
import java.util.TreeMap;
3634

3735
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
3836
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
@@ -44,33 +42,25 @@ public class DataFrameTransformState {
4442
private static final ParseField CURRENT_POSITION = new ParseField("current_position");
4543
private static final ParseField CHECKPOINT = new ParseField("checkpoint");
4644
private static final ParseField REASON = new ParseField("reason");
45+
private static final ParseField PROGRESS = new ParseField("progress");
4746

4847
@SuppressWarnings("unchecked")
4948
public static final ConstructingObjectParser<DataFrameTransformState, Void> PARSER =
5049
new ConstructingObjectParser<>("data_frame_transform_state", true,
5150
args -> new DataFrameTransformState((DataFrameTransformTaskState) args[0],
5251
(IndexerState) args[1],
53-
(HashMap<String, Object>) args[2],
52+
(Map<String, Object>) args[2],
5453
(long) args[3],
55-
(String) args[4]));
54+
(String) args[4],
55+
(DataFrameTransformProgress) args[5]));
5656

5757
static {
58-
PARSER.declareField(constructorArg(),
59-
p -> DataFrameTransformTaskState.fromString(p.text()),
60-
TASK_STATE,
61-
ObjectParser.ValueType.STRING);
62-
PARSER.declareField(constructorArg(), p -> IndexerState.fromString(p.text()), INDEXER_STATE, ObjectParser.ValueType.STRING);
63-
PARSER.declareField(optionalConstructorArg(), p -> {
64-
if (p.currentToken() == XContentParser.Token.START_OBJECT) {
65-
return p.map();
66-
}
67-
if (p.currentToken() == XContentParser.Token.VALUE_NULL) {
68-
return null;
69-
}
70-
throw new IllegalArgumentException("Unsupported token [" + p.currentToken() + "]");
71-
}, CURRENT_POSITION, ObjectParser.ValueType.VALUE_OBJECT_ARRAY);
58+
PARSER.declareField(constructorArg(), p -> DataFrameTransformTaskState.fromString(p.text()), TASK_STATE, ValueType.STRING);
59+
PARSER.declareField(constructorArg(), p -> IndexerState.fromString(p.text()), INDEXER_STATE, ValueType.STRING);
60+
PARSER.declareField(optionalConstructorArg(), (p, c) -> p.mapOrdered(), CURRENT_POSITION, ValueType.OBJECT);
7261
PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), CHECKPOINT);
7362
PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), REASON);
63+
PARSER.declareField(optionalConstructorArg(), DataFrameTransformProgress::fromXContent, PROGRESS, ValueType.OBJECT);
7464
}
7565

7666
public static DataFrameTransformState fromXContent(XContentParser parser) throws IOException {
@@ -80,19 +70,22 @@ public static DataFrameTransformState fromXContent(XContentParser parser) throws
8070
private final DataFrameTransformTaskState taskState;
8171
private final IndexerState indexerState;
8272
private final long checkpoint;
83-
private final SortedMap<String, Object> currentPosition;
73+
private final Map<String, Object> currentPosition;
8474
private final String reason;
75+
private final DataFrameTransformProgress progress;
8576

8677
public DataFrameTransformState(DataFrameTransformTaskState taskState,
8778
IndexerState indexerState,
8879
@Nullable Map<String, Object> position,
8980
long checkpoint,
90-
@Nullable String reason) {
81+
@Nullable String reason,
82+
@Nullable DataFrameTransformProgress progress) {
9183
this.taskState = taskState;
9284
this.indexerState = indexerState;
93-
this.currentPosition = position == null ? null : Collections.unmodifiableSortedMap(new TreeMap<>(position));
85+
this.currentPosition = position == null ? null : Collections.unmodifiableMap(new LinkedHashMap<>(position));
9486
this.checkpoint = checkpoint;
9587
this.reason = reason;
88+
this.progress = progress;
9689
}
9790

9891
public IndexerState getIndexerState() {
@@ -117,6 +110,11 @@ public String getReason() {
117110
return reason;
118111
}
119112

113+
@Nullable
114+
public DataFrameTransformProgress getProgress() {
115+
return progress;
116+
}
117+
120118
@Override
121119
public boolean equals(Object other) {
122120
if (this == other) {
@@ -132,13 +130,14 @@ public boolean equals(Object other) {
132130
return Objects.equals(this.taskState, that.taskState) &&
133131
Objects.equals(this.indexerState, that.indexerState) &&
134132
Objects.equals(this.currentPosition, that.currentPosition) &&
133+
Objects.equals(this.progress, that.progress) &&
135134
this.checkpoint == that.checkpoint &&
136135
Objects.equals(this.reason, that.reason);
137136
}
138137

139138
@Override
140139
public int hashCode() {
141-
return Objects.hash(taskState, indexerState, currentPosition, checkpoint, reason);
140+
return Objects.hash(taskState, indexerState, currentPosition, checkpoint, reason, progress);
142141
}
143142

144143
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ public static DataFrameTransformStateAndStats fromXContent(XContentParser parser
5757
private final DataFrameTransformCheckpointingInfo checkpointingInfo;
5858

5959
public DataFrameTransformStateAndStats(String id, DataFrameTransformState state, DataFrameIndexerTransformStats stats,
60-
DataFrameTransformCheckpointingInfo checkpointingInfo) {
60+
DataFrameTransformCheckpointingInfo checkpointingInfo) {
6161
this.id = id;
6262
this.transformState = state;
6363
this.transformStats = stats;

‎client/rest-high-level/src/test/java/org/elasticsearch/client/DataFrameTransformIT.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
import static org.hamcrest.Matchers.containsString;
7272
import static org.hamcrest.Matchers.empty;
7373
import static org.hamcrest.Matchers.equalTo;
74+
import static org.hamcrest.Matchers.greaterThan;
7475
import static org.hamcrest.Matchers.hasSize;
7576
import static org.hamcrest.Matchers.is;
7677

@@ -360,6 +361,10 @@ public void testGetStats() throws Exception {
360361
assertEquals(DataFrameTransformTaskState.STARTED, stateAndStats.getTransformState().getTaskState());
361362
assertEquals(null, stateAndStats.getTransformState().getReason());
362363
assertNotEquals(zeroIndexerStats, stateAndStats.getTransformStats());
364+
assertNotNull(stateAndStats.getTransformState().getProgress());
365+
assertThat(stateAndStats.getTransformState().getProgress().getPercentComplete(), equalTo(100.0));
366+
assertThat(stateAndStats.getTransformState().getProgress().getTotalDocs(), greaterThan(0L));
367+
assertThat(stateAndStats.getTransformState().getProgress().getRemainingDocs(), equalTo(0L));
363368
});
364369
}
365370
}

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
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+
20+
package org.elasticsearch.client.dataframe.transforms;
21+
22+
import org.elasticsearch.common.xcontent.XContentBuilder;
23+
import org.elasticsearch.test.ESTestCase;
24+
25+
import java.io.IOException;
26+
27+
import static org.elasticsearch.test.AbstractXContentTestCase.xContentTester;
28+
29+
public class DataFrameTransformProgressTests extends ESTestCase {
30+
31+
public void testFromXContent() throws IOException {
32+
xContentTester(this::createParser,
33+
DataFrameTransformProgressTests::randomInstance,
34+
DataFrameTransformProgressTests::toXContent,
35+
DataFrameTransformProgress::fromXContent)
36+
.supportsUnknownFields(true)
37+
.randomFieldsExcludeFilter(field -> field.startsWith("state"))
38+
.test();
39+
}
40+
41+
public static DataFrameTransformProgress randomInstance() {
42+
long totalDocs = randomNonNegativeLong();
43+
Long docsRemaining = randomBoolean() ? null : randomLongBetween(0, totalDocs);
44+
double percentComplete = totalDocs == 0 ? 1.0 : docsRemaining == null ? 0.0 : 100.0*(double)(totalDocs - docsRemaining)/totalDocs;
45+
return new DataFrameTransformProgress(totalDocs, docsRemaining, percentComplete);
46+
}
47+
48+
public static void toXContent(DataFrameTransformProgress progress, XContentBuilder builder) throws IOException {
49+
builder.startObject();
50+
builder.field(DataFrameTransformProgress.TOTAL_DOCS.getPreferredName(), progress.getTotalDocs());
51+
builder.field(DataFrameTransformProgress.DOCS_REMAINING.getPreferredName(), progress.getRemainingDocs());
52+
builder.field(DataFrameTransformProgress.PERCENT_COMPLETE.getPreferredName(), progress.getPercentComplete());
53+
builder.endObject();
54+
}
55+
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ public void testFromXContent() throws IOException {
4040

4141
public static DataFrameTransformStateAndStats randomInstance() {
4242
return new DataFrameTransformStateAndStats(randomAlphaOfLength(10),
43-
DataFrameTransformStateTests.randomDataFrameTransformState(),
44-
DataFrameIndexerTransformStatsTests.randomStats(),
45-
DataFrameTransformCheckpointingInfoTests.randomDataFrameTransformCheckpointingInfo());
43+
DataFrameTransformStateTests.randomDataFrameTransformState(),
44+
DataFrameIndexerTransformStatsTests.randomStats(),
45+
DataFrameTransformCheckpointingInfoTests.randomDataFrameTransformCheckpointingInfo());
4646
}
4747

4848
public static void toXContent(DataFrameTransformStateAndStats stateAndStats, XContentBuilder builder) throws IOException {

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import org.elasticsearch.test.ESTestCase;
2525

2626
import java.io.IOException;
27-
import java.util.HashMap;
27+
import java.util.LinkedHashMap;
2828
import java.util.Map;
2929

3030
import static org.elasticsearch.test.AbstractXContentTestCase.xContentTester;
@@ -46,7 +46,8 @@ public static DataFrameTransformState randomDataFrameTransformState() {
4646
randomFrom(IndexerState.values()),
4747
randomPositionMap(),
4848
randomLongBetween(0,10),
49-
randomBoolean() ? null : randomAlphaOfLength(10));
49+
randomBoolean() ? null : randomAlphaOfLength(10),
50+
randomBoolean() ? null : DataFrameTransformProgressTests.randomInstance());
5051
}
5152

5253
public static void toXContent(DataFrameTransformState state, XContentBuilder builder) throws IOException {
@@ -60,6 +61,10 @@ public static void toXContent(DataFrameTransformState state, XContentBuilder bui
6061
if (state.getReason() != null) {
6162
builder.field("reason", state.getReason());
6263
}
64+
if (state.getProgress() != null) {
65+
builder.field("progress");
66+
DataFrameTransformProgressTests.toXContent(state.getProgress(), builder);
67+
}
6368
builder.endObject();
6469
}
6570

@@ -68,7 +73,7 @@ private static Map<String, Object> randomPositionMap() {
6873
return null;
6974
}
7075
int numFields = randomIntBetween(1, 5);
71-
Map<String, Object> position = new HashMap<>();
76+
Map<String, Object> position = new LinkedHashMap<>();
7277
for (int i = 0; i < numFields; i++) {
7378
Object value;
7479
if (randomBoolean()) {

0 commit comments

Comments
 (0)