You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ja...@apache.org on 2023/09/07 00:50:56 UTC
[iotdb] branch master updated: [IoTDB ML] Basic model training and inference features available
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new d33b16b1ffd [IoTDB ML] Basic model training and inference features available
d33b16b1ffd is described below
commit d33b16b1ffdec5d1c07fd3ad54f24963f038a079
Author: YangCaiyin <yc...@gmail.com>
AuthorDate: Thu Sep 7 08:50:50 2023 +0800
[IoTDB ML] Basic model training and inference features available
---
.../java/org/apache/iotdb/rpc/TSStatusCode.java | 1 +
.../org/apache/iotdb/db/qp/sql/IdentifierParser.g4 | 7 +-
.../org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 | 35 +-
.../antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 | 20 +-
.../consensus/request/ConfigPhysicalPlan.java | 10 +-
.../consensus/request/ConfigPhysicalPlanType.java | 3 +-
.../{ShowTrailPlan.java => GetModelInfoPlan.java} | 38 +-
.../{ShowTrailPlan.java => ShowTrialPlan.java} | 38 +-
.../request/write/model/UpdateModelInfoPlan.java | 16 +-
.../request/write/model/UpdateModelStatePlan.java | 18 +-
.../response/model/GetModelInfoResp.java} | 31 +-
.../response/{ => model}/ModelTableResp.java | 2 +-
.../TrialTableResp.java} | 28 +-
.../iotdb/confignode/manager/ConfigManager.java | 20 +-
.../apache/iotdb/confignode/manager/IManager.java | 13 +-
.../iotdb/confignode/manager/ModelManager.java | 99 +++--
.../iotdb/confignode/manager/ProcedureManager.java | 5 +-
.../iotdb/confignode/persistence/ModelInfo.java | 147 +++++---
.../persistence/executor/ConfigPlanExecutor.java | 9 +-
.../procedure/impl/model/CreateModelProcedure.java | 125 +++----
.../procedure/impl/model/DropModelProcedure.java | 131 +++----
.../procedure/state/model/CreateModelState.java | 2 -
.../thrift/ConfigNodeRPCServiceProcessor.java | 15 +-
.../iotdb/db/protocol/client/ConfigNodeClient.java | 26 +-
.../protocol/thrift/impl/MLNodeRPCServiceImpl.java | 14 +-
.../common/header/ColumnHeaderConstant.java | 9 +-
.../common/header/DatasetHeaderFactory.java | 4 +-
.../fragment/FragmentInstanceManager.java | 11 +
.../execution/operator/AggregationUtil.java | 2 +-
.../operator/process/ml/ForecastOperator.java | 241 ++++++++++++
.../db/queryengine/plan/analyze/Analysis.java | 11 +
.../queryengine/plan/analyze/AnalyzeVisitor.java | 118 +++++-
.../plan/analyze/ExpressionAnalyzer.java | 7 +
.../plan/analyze/ExpressionTypeAnalyzer.java | 2 +
.../queryengine/plan/analyze/IModelFetcher.java} | 12 +-
.../db/queryengine/plan/analyze/ModelFetcher.java | 75 ++++
.../plan/execution/config/ConfigTaskVisitor.java | 8 +-
.../config/executor/ClusterConfigTaskExecutor.java | 259 ++++++++-----
.../config/executor/IConfigTaskExecutor.java | 2 +-
.../config/metadata/model/ShowModelsTask.java | 12 +-
.../{ShowTrailsTask.java => ShowTrialsTask.java} | 10 +-
.../plan/expression/ExpressionFactory.java | 9 +
.../plan/expression/multi/FunctionExpression.java | 10 +
.../plan/expression/multi/FunctionType.java | 3 +-
.../visitor/RemoveRootPrefixVisitor.java | 45 +++
.../db/queryengine/plan/parser/ASTVisitor.java | 42 ++-
.../plan/parser/StatementGenerator.java | 50 +--
.../plan/planner/LogicalPlanBuilder.java | 10 +
.../plan/planner/LogicalPlanVisitor.java | 20 +
.../plan/planner/OperatorTreeGenerator.java | 49 +++
.../plan/planner/plan/node/PlanGraphPrinter.java | 12 +
.../plan/planner/plan/node/PlanNodeType.java | 5 +-
.../plan/planner/plan/node/PlanVisitor.java | 5 +
.../planner/plan/node/process/ml/ForecastNode.java | 122 ++++++
.../model/ForecastModelInferenceDescriptor.java | 176 +++++++++
.../parameter/model/ModelInferenceDescriptor.java | 111 ++++++
.../plan/statement/StatementVisitor.java | 6 +-
.../plan/statement/component/ResultColumn.java | 5 +-
.../plan/statement/component/SelectComponent.java | 10 +-
.../plan/statement/crud/QueryStatement.java | 52 ++-
.../metadata/model/CreateModelStatement.java | 71 ++--
...ailsStatement.java => ShowTrialsStatement.java} | 6 +-
.../apache/iotdb/db/service/MLNodeRPCService.java | 8 +-
.../exception/ModelInferenceProcessException.java} | 12 +-
.../iotdb/db/utils/TimestampPrecisionUtils.java | 10 +-
.../iotdb/db/utils/constant/SqlConstant.java | 3 +
.../plan/parser/StatementGeneratorTest.java | 9 +-
iotdb-core/mlnode/.gitignore | 3 +-
iotdb-core/mlnode/README.md | 4 +-
.../mlnode/iotdb/mlnode/algorithm/factory.py | 66 ++++
.../iotdb/mlnode/algorithm/hyperparameter.py | 407 +++++++++++++++++++++
iotdb-core/mlnode/iotdb/mlnode/algorithm/metric.py | 80 ++++
.../mlnode/algorithm/models/forecast/__init__.py} | 10 -
.../mlnode/algorithm/models/forecast/dlinear.py | 133 +++++++
.../mlnode/algorithm/models/forecast/nbeats.py | 131 +++++++
.../mlnode/iotdb/mlnode/algorithm/validator.py | 45 +++
iotdb-core/mlnode/iotdb/mlnode/client.py | 111 ++----
iotdb-core/mlnode/iotdb/mlnode/config.py | 36 ++
iotdb-core/mlnode/iotdb/mlnode/constant.py | 68 +++-
iotdb-core/mlnode/iotdb/mlnode/dataset/dataset.py | 74 ++++
iotdb-core/mlnode/iotdb/mlnode/dataset/factory.py | 50 +++
iotdb-core/mlnode/iotdb/mlnode/dataset/source.py | 80 ++++
.../mlnode/dataset/utils/__init__.py} | 10 -
.../iotdb/mlnode/dataset/utils/time_features.py | 149 ++++++++
iotdb-core/mlnode/iotdb/mlnode/exception.py | 39 +-
iotdb-core/mlnode/iotdb/mlnode/handler.py | 63 +++-
iotdb-core/mlnode/iotdb/mlnode/parser.py | 99 +++++
.../mlnode/process/__init__.py} | 10 -
iotdb-core/mlnode/iotdb/mlnode/process/manager.py | 124 +++++++
iotdb-core/mlnode/iotdb/mlnode/process/task.py | 265 ++++++++++++++
iotdb-core/mlnode/iotdb/mlnode/process/trial.py | 224 ++++++++++++
iotdb-core/mlnode/iotdb/mlnode/serde.py | 147 ++++++--
iotdb-core/mlnode/iotdb/mlnode/service.py | 4 -
iotdb-core/mlnode/iotdb/mlnode/storage.py | 28 +-
iotdb-core/mlnode/iotdb/mlnode/util.py | 25 +-
iotdb-core/mlnode/requirements.txt | 7 +-
iotdb-core/mlnode/resources/conf/iotdb-mlnode.toml | 14 +-
.../mlnode/test/test_create_forecast_model.py | 100 +++++
iotdb-core/mlnode/test/test_model_storage.py | 17 +-
iotdb-core/mlnode/test/test_serde.py | 8 +-
.../resources/conf/iotdb-common.properties | 6 +
.../iotdb/commons/client/ClientPoolFactory.java | 26 ++
.../iotdb/commons/client/mlnode}/MLNodeClient.java | 144 ++++++--
.../client/mlnode/MLNodeClientManager.java} | 25 +-
.../iotdb/commons/client/mlnode/MLNodeInfo.java} | 14 +-
.../apache/iotdb/commons/conf/CommonConfig.java | 11 +
.../iotdb/commons/conf/CommonDescriptor.java | 5 +
.../commons/model/ForecastModeInformation.java | 153 ++++++++
.../iotdb/commons/model/ModelInformation.java | 235 ++++++------
...TrailInformation.java => TrialInformation.java} | 34 +-
.../udf/builtin/ModelInferenceFunction.java} | 34 +-
.../commons/udf/service/UDFManagementService.java | 19 +
.../tsfile/file/metadata/enums/TSDataType.java | 10 +
.../iotdb/tsfile/read/common/block/TsBlock.java | 20 +-
.../tsfile/read/common/block/TsBlockBuilder.java | 4 +-
.../thrift-commons/src/main/thrift/common.thrift | 5 +-
.../src/main/thrift/confignode.thrift | 40 +-
.../src/main/thrift/datanode.thrift | 7 +-
.../thrift-mlnode/src/main/thrift/mlnode.thrift | 12 +-
119 files changed, 4921 insertions(+), 1011 deletions(-)
diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
index 173fd7c8899..bd518452b4b 100644
--- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
+++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
@@ -189,6 +189,7 @@ public enum TSStatusCode {
// ML Model
CREATE_MODEL_ERROR(1500),
DROP_MODEL_ERROR(1501),
+ MODEL_EXIST_ERROR(1502),
MLNODE_INTERNAL_ERROR(1510),
// Pipe Plugin
diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4
index 95aea43539e..79b953e5b1b 100644
--- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4
+++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4
@@ -45,7 +45,6 @@ keyWords
| AS
| ASC
| ATTRIBUTES
- | AUTO
| BEFORE
| BEGIN
| BETWEEN
@@ -77,6 +76,7 @@ keyWords
| DATABASES
| DATANODEID
| DATANODES
+ | DATASET
| DEACTIVATE
| DEBUG
| DELETE
@@ -84,7 +84,6 @@ keyWords
| DESCRIBE
| DETAILS
| DEVICE
- | DEVICEID
| DEVICES
| DISABLE
| DISCARD
@@ -110,6 +109,7 @@ keyWords
| GRANT
| GROUP
| HAVING
+ | HYPERPARAMETERS
| IN
| INDEX
| INFO
@@ -146,6 +146,7 @@ keyWords
| OFF
| OFFSET
| ON
+ | OPTIONS
| OR
| ORDER
| ONSUCCESS
@@ -224,7 +225,7 @@ keyWords
| TOLERANCE
| TOP
| TRACING
- | TRAILS
+ | TRIALS
| TRIGGER
| TRIGGERS
| TRUE
diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4
index c9d66ee181e..6a447bcb1e1 100644
--- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4
+++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4
@@ -62,7 +62,7 @@ ddlStatement
| showVariables | showCluster | showRegions | showDataNodes | showConfigNodes
| getRegionId | getTimeSlotList | countTimeSlotList | getSeriesSlotList | migrateRegion
// ML Model
- | createModel | dropModel | showModels | showTrails
+ | createModel | dropModel | showModels | showTrials
// Quota
| setSpaceQuota | showSpaceQuota | setThrottleQuota | showThrottleQuota
// View
@@ -585,11 +585,28 @@ showPipePlugins
// ML Model =========================================================================================
// ---- Create Model
createModel
- : CREATE AUTO? MODEL modelId=identifier
- WITH attributePair (COMMA attributePair)*
- BEGIN
- selectStatement
- END
+ : CREATE MODEL modelId=identifier
+ OPTIONS LR_BRACKET attributePair (COMMA attributePair)* RR_BRACKET
+ (WITH HYPERPARAMETERS LR_BRACKET hparamPair (COMMA hparamPair)* RR_BRACKET)?
+ ON DATASET LR_BRACKET selectStatement RR_BRACKET
+ ;
+
+hparamPair
+ : hparamKey=attributeKey operator_eq hparamValue
+ ;
+
+hparamValue
+ : attributeValue
+ | hparamRange
+ | hparamCandidates
+ ;
+
+hparamRange
+ : LR_BRACKET hparamRangeStart=attributeValue COMMA hparamRangeEnd=attributeValue RR_BRACKET
+ ;
+
+hparamCandidates
+ : LS_BRACKET attributeValue (COMMA attributeValue)* RS_BRACKET
;
// ---- Drop Model
@@ -602,9 +619,9 @@ showModels
: SHOW MODELS
;
-// ---- Show Trails
-showTrails
- : SHOW TRAILS modelId=identifier
+// ---- Show Trials
+showTrials
+ : SHOW TRIALS modelId=identifier
;
// Create Logical View
diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4
index 26e62b1ec49..41baa980425 100644
--- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4
+++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4
@@ -97,10 +97,6 @@ ATTRIBUTES
: A T T R I B U T E S
;
-AUTO
- : A U T O
- ;
-
BEFORE
: B E F O R E
;
@@ -214,6 +210,10 @@ DATANODES
: D A T A N O D E S
;
+DATASET
+ : D A T A S E T
+ ;
+
DEACTIVATE
: D E A C T I V A T E
;
@@ -338,6 +338,10 @@ HAVING
: H A V I N G
;
+HYPERPARAMETERS
+ : H Y P E R P A R A M E T E R S
+ ;
+
IN
: I N
;
@@ -482,6 +486,10 @@ ON
: O N
;
+OPTIONS
+ : O P T I O N S
+ ;
+
OR
: O R
;
@@ -786,8 +794,8 @@ TRACING
: T R A C I N G
;
-TRAILS
- : T R A I L S
+TRIALS
+ : T R I A L S
;
TRIGGER
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java
index b097bdb7499..4d1f5e0cc96 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java
@@ -26,8 +26,9 @@ import org.apache.iotdb.confignode.consensus.request.read.database.GetDatabasePl
import org.apache.iotdb.confignode.consensus.request.read.datanode.GetDataNodeConfigurationPlan;
import org.apache.iotdb.confignode.consensus.request.read.function.GetFunctionTablePlan;
import org.apache.iotdb.confignode.consensus.request.read.function.GetUDFJarPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.GetModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.read.model.ShowModelPlan;
-import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrailPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrialPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.CountTimeSlotListPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.GetDataPartitionPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.GetNodePathsPartitionPlan;
@@ -461,8 +462,11 @@ public abstract class ConfigPhysicalPlan implements IConsensusRequest {
case ShowModel:
plan = new ShowModelPlan();
break;
- case ShowTrail:
- plan = new ShowTrailPlan();
+ case ShowTrial:
+ plan = new ShowTrialPlan();
+ break;
+ case GetModelInfo:
+ plan = new GetModelInfoPlan();
break;
case CreatePipePlugin:
plan = new CreatePipePluginPlan();
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java
index 3392b283646..2afdb80d82e 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java
@@ -162,7 +162,8 @@ public enum ConfigPhysicalPlanType {
UpdateModelState((short) 1202),
DropModel((short) 1203),
ShowModel((short) 1204),
- ShowTrail((short) 1205),
+ ShowTrial((short) 1205),
+ GetModelInfo((short) 1206),
/** Pipe Plugin. */
CreatePipePlugin((short) 1300),
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/GetModelInfoPlan.java
similarity index 65%
copy from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java
copy to iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/GetModelInfoPlan.java
index 6a982585e84..79c7e9e0142 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/GetModelInfoPlan.java
@@ -21,7 +21,7 @@ package org.apache.iotdb.confignode.consensus.request.read.model;
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlan;
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
import java.io.DataOutputStream;
@@ -29,50 +29,32 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Objects;
-public class ShowTrailPlan extends ConfigPhysicalPlan {
+public class GetModelInfoPlan extends ConfigPhysicalPlan {
private String modelId;
- private String trailId;
- public ShowTrailPlan() {
- super(ConfigPhysicalPlanType.ShowTrail);
+ public GetModelInfoPlan() {
+ super(ConfigPhysicalPlanType.GetModelInfo);
}
- public ShowTrailPlan(TShowTrailReq showTrailReq) {
- super(ConfigPhysicalPlanType.ShowTrail);
- this.modelId = showTrailReq.getModelId();
- if (showTrailReq.isSetTrailId()) {
- this.trailId = showTrailReq.getTrailId();
- }
+ public GetModelInfoPlan(TGetModelInfoReq getModelInfoReq) {
+ super(ConfigPhysicalPlanType.GetModelInfo);
+ this.modelId = getModelInfoReq.getModelId();
}
public String getModelId() {
return modelId;
}
- public boolean isSetTrailId() {
- return trailId != null;
- }
-
- public String getTrailId() {
- return trailId;
- }
-
@Override
protected void serializeImpl(DataOutputStream stream) throws IOException {
stream.writeShort(getType().getPlanType());
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(trailId != null, stream);
- ReadWriteIOUtils.write(trailId, stream);
}
@Override
protected void deserializeImpl(ByteBuffer buffer) throws IOException {
this.modelId = ReadWriteIOUtils.readString(buffer);
- boolean isSetTrailId = ReadWriteIOUtils.readBool(buffer);
- if (isSetTrailId) {
- this.trailId = ReadWriteIOUtils.readString(buffer);
- }
}
@Override
@@ -86,12 +68,12 @@ public class ShowTrailPlan extends ConfigPhysicalPlan {
if (!super.equals(o)) {
return false;
}
- ShowTrailPlan that = (ShowTrailPlan) o;
- return modelId.equals(that.modelId) && Objects.equals(trailId, that.trailId);
+ GetModelInfoPlan that = (GetModelInfoPlan) o;
+ return Objects.equals(modelId, that.modelId);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), modelId, trailId);
+ return Objects.hash(super.hashCode(), modelId);
}
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrialPlan.java
similarity index 70%
rename from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java
rename to iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrialPlan.java
index 6a982585e84..136e1c2408d 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrailPlan.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowTrialPlan.java
@@ -21,7 +21,7 @@ package org.apache.iotdb.confignode.consensus.request.read.model;
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlan;
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
import java.io.DataOutputStream;
@@ -29,20 +29,20 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Objects;
-public class ShowTrailPlan extends ConfigPhysicalPlan {
+public class ShowTrialPlan extends ConfigPhysicalPlan {
private String modelId;
- private String trailId;
+ private String trialId;
- public ShowTrailPlan() {
- super(ConfigPhysicalPlanType.ShowTrail);
+ public ShowTrialPlan() {
+ super(ConfigPhysicalPlanType.ShowTrial);
}
- public ShowTrailPlan(TShowTrailReq showTrailReq) {
- super(ConfigPhysicalPlanType.ShowTrail);
+ public ShowTrialPlan(TShowTrialReq showTrailReq) {
+ super(ConfigPhysicalPlanType.ShowTrial);
this.modelId = showTrailReq.getModelId();
- if (showTrailReq.isSetTrailId()) {
- this.trailId = showTrailReq.getTrailId();
+ if (showTrailReq.isSetTrialId()) {
+ this.trialId = showTrailReq.getTrialId();
}
}
@@ -50,20 +50,20 @@ public class ShowTrailPlan extends ConfigPhysicalPlan {
return modelId;
}
- public boolean isSetTrailId() {
- return trailId != null;
+ public boolean isSetTrialId() {
+ return trialId != null;
}
- public String getTrailId() {
- return trailId;
+ public String getTrialId() {
+ return trialId;
}
@Override
protected void serializeImpl(DataOutputStream stream) throws IOException {
stream.writeShort(getType().getPlanType());
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(trailId != null, stream);
- ReadWriteIOUtils.write(trailId, stream);
+ ReadWriteIOUtils.write(trialId != null, stream);
+ ReadWriteIOUtils.write(trialId, stream);
}
@Override
@@ -71,7 +71,7 @@ public class ShowTrailPlan extends ConfigPhysicalPlan {
this.modelId = ReadWriteIOUtils.readString(buffer);
boolean isSetTrailId = ReadWriteIOUtils.readBool(buffer);
if (isSetTrailId) {
- this.trailId = ReadWriteIOUtils.readString(buffer);
+ this.trialId = ReadWriteIOUtils.readString(buffer);
}
}
@@ -86,12 +86,12 @@ public class ShowTrailPlan extends ConfigPhysicalPlan {
if (!super.equals(o)) {
return false;
}
- ShowTrailPlan that = (ShowTrailPlan) o;
- return modelId.equals(that.modelId) && Objects.equals(trailId, that.trailId);
+ ShowTrialPlan that = (ShowTrialPlan) o;
+ return modelId.equals(that.modelId) && Objects.equals(trialId, that.trialId);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), modelId, trailId);
+ return Objects.hash(super.hashCode(), modelId, trialId);
}
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelInfoPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelInfoPlan.java
index 52c1100a00c..5261472668d 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelInfoPlan.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelInfoPlan.java
@@ -33,7 +33,7 @@ import java.util.Objects;
public class UpdateModelInfoPlan extends ConfigPhysicalPlan {
private String modelId;
- private String trailId;
+ private String trialId;
private Map<String, String> modelInfo;
public UpdateModelInfoPlan() {
@@ -43,7 +43,7 @@ public class UpdateModelInfoPlan extends ConfigPhysicalPlan {
public UpdateModelInfoPlan(TUpdateModelInfoReq updateModelInfoReq) {
super(ConfigPhysicalPlanType.UpdateModelInfo);
this.modelId = updateModelInfoReq.getModelId();
- this.trailId = updateModelInfoReq.getTrailId();
+ this.trialId = updateModelInfoReq.getTrialId();
this.modelInfo = updateModelInfoReq.getModelInfo();
}
@@ -51,8 +51,8 @@ public class UpdateModelInfoPlan extends ConfigPhysicalPlan {
return modelId;
}
- public String getTrailId() {
- return trailId;
+ public String getTrialId() {
+ return trialId;
}
public Map<String, String> getModelInfo() {
@@ -63,14 +63,14 @@ public class UpdateModelInfoPlan extends ConfigPhysicalPlan {
protected void serializeImpl(DataOutputStream stream) throws IOException {
stream.writeShort(getType().getPlanType());
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(trailId, stream);
+ ReadWriteIOUtils.write(trialId, stream);
ReadWriteIOUtils.write(modelInfo, stream);
}
@Override
protected void deserializeImpl(ByteBuffer buffer) throws IOException {
this.modelId = ReadWriteIOUtils.readString(buffer);
- this.trailId = ReadWriteIOUtils.readString(buffer);
+ this.trialId = ReadWriteIOUtils.readString(buffer);
this.modelInfo = ReadWriteIOUtils.readMap(buffer);
}
@@ -87,12 +87,12 @@ public class UpdateModelInfoPlan extends ConfigPhysicalPlan {
}
UpdateModelInfoPlan that = (UpdateModelInfoPlan) o;
return modelId.equals(that.modelId)
- && trailId.equals(that.trailId)
+ && trialId.equals(that.trialId)
&& modelInfo.equals(that.modelInfo);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), modelId, trailId, modelInfo);
+ return Objects.hash(super.hashCode(), modelId, trialId, modelInfo);
}
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelStatePlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelStatePlan.java
index 06a88e4e272..6ac3fc1a5c1 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelStatePlan.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/model/UpdateModelStatePlan.java
@@ -34,7 +34,7 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
private String modelId;
private TrainingState state;
- private String bestTrailId;
+ private String bestTrialId;
public UpdateModelStatePlan() {
super(ConfigPhysicalPlanType.UpdateModelState);
@@ -44,7 +44,7 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
super(ConfigPhysicalPlanType.UpdateModelState);
this.modelId = updateModelStateReq.getModelId();
this.state = updateModelStateReq.getState();
- this.bestTrailId = updateModelStateReq.getBestTrailId();
+ this.bestTrialId = updateModelStateReq.getBestTrialId();
}
public String getModelId() {
@@ -55,8 +55,8 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
return state;
}
- public String getBestTrailId() {
- return bestTrailId;
+ public String getBestTrialId() {
+ return bestTrialId;
}
@Override
@@ -64,10 +64,10 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
stream.writeShort(getType().getPlanType());
ReadWriteIOUtils.write(modelId, stream);
ReadWriteIOUtils.write(state.getValue(), stream);
- boolean isNull = bestTrailId == null;
+ boolean isNull = bestTrialId == null;
ReadWriteIOUtils.write(isNull, stream);
if (!isNull) {
- ReadWriteIOUtils.write(bestTrailId, stream);
+ ReadWriteIOUtils.write(bestTrialId, stream);
}
}
@@ -77,7 +77,7 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
this.state = TrainingState.findByValue(ReadWriteIOUtils.readInt(buffer));
boolean isNull = ReadWriteIOUtils.readBool(buffer);
if (!isNull) {
- this.bestTrailId = ReadWriteIOUtils.readString(buffer);
+ this.bestTrialId = ReadWriteIOUtils.readString(buffer);
}
}
@@ -95,11 +95,11 @@ public class UpdateModelStatePlan extends ConfigPhysicalPlan {
UpdateModelStatePlan that = (UpdateModelStatePlan) o;
return modelId.equals(that.modelId)
&& state == that.state
- && Objects.equals(bestTrailId, that.bestTrailId);
+ && Objects.equals(bestTrialId, that.bestTrialId);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), modelId, state, bestTrailId);
+ return Objects.hash(super.hashCode(), modelId, state, bestTrialId);
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/GetModelInfoResp.java
similarity index 50%
copy from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
copy to iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/GetModelInfoResp.java
index 336bbcce0bb..dd4a842d6ec 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/GetModelInfoResp.java
@@ -17,11 +17,30 @@
* under the License.
*/
-package org.apache.iotdb.db.queryengine.plan.expression.multi;
+package org.apache.iotdb.confignode.consensus.response.model;
-/** */
-public enum FunctionType {
- AGGREGATION_FUNCTION,
- BUILT_IN_SCALAR_FUNCTION,
- UDF
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
+import org.apache.iotdb.consensus.common.DataSet;
+
+import java.nio.ByteBuffer;
+
+public class GetModelInfoResp implements DataSet {
+
+ private final TSStatus status;
+ private ByteBuffer serializedModelInformation;
+
+ public GetModelInfoResp(TSStatus status) {
+ this.status = status;
+ }
+
+ public void setModelInfo(ByteBuffer serializedModelInformation) {
+ this.serializedModelInformation = serializedModelInformation;
+ }
+
+ public TGetModelInfoResp convertToThriftResponse() {
+ TGetModelInfoResp resp = new TGetModelInfoResp(status);
+ resp.setModelInfo(serializedModelInformation);
+ return resp;
+ }
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/ModelTableResp.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java
similarity index 97%
rename from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/ModelTableResp.java
rename to iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java
index 6642f76be97..9a23d9ed713 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/ModelTableResp.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java
@@ -17,7 +17,7 @@
* under the License.
*/
-package org.apache.iotdb.confignode.consensus.response;
+package org.apache.iotdb.confignode.consensus.response.model;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.model.ModelInformation;
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/TrailTableResp.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/TrialTableResp.java
similarity index 57%
rename from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/TrailTableResp.java
rename to iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/TrialTableResp.java
index 1f9a6b5acbf..279703efa11 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/TrailTableResp.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/TrialTableResp.java
@@ -17,11 +17,11 @@
* under the License.
*/
-package org.apache.iotdb.confignode.consensus.response;
+package org.apache.iotdb.confignode.consensus.response.model;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
-import org.apache.iotdb.commons.model.TrailInformation;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.commons.model.TrialInformation;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.consensus.common.DataSet;
import java.io.IOException;
@@ -29,27 +29,27 @@ import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
-public class TrailTableResp implements DataSet {
+public class TrialTableResp implements DataSet {
private final TSStatus status;
- private final List<ByteBuffer> serializedAllTrailInformation;
+ private final List<ByteBuffer> serializedAllTrialInformation;
- public TrailTableResp(TSStatus status) {
+ public TrialTableResp(TSStatus status) {
this.status = status;
- this.serializedAllTrailInformation = new ArrayList<>();
+ this.serializedAllTrialInformation = new ArrayList<>();
}
- public void addTrailInformation(TrailInformation trailInformation) throws IOException {
- this.serializedAllTrailInformation.add(trailInformation.serializeShowTrailResult());
+ public void addTrialInformation(TrialInformation trialInformation) throws IOException {
+ this.serializedAllTrialInformation.add(trialInformation.serializeShowTrialResult());
}
- public void addTrailInformation(List<TrailInformation> trailInformationList) throws IOException {
- for (TrailInformation trailInformation : trailInformationList) {
- this.serializedAllTrailInformation.add(trailInformation.serializeShowTrailResult());
+ public void addTrialInformation(List<TrialInformation> trialInformationList) throws IOException {
+ for (TrialInformation trialInformation : trialInformationList) {
+ this.serializedAllTrialInformation.add(trialInformation.serializeShowTrialResult());
}
}
- public TShowTrailResp convertToThriftResponse() throws IOException {
- return new TShowTrailResp(status, serializedAllTrailInformation);
+ public TShowTrialResp convertToThriftResponse() throws IOException {
+ return new TShowTrialResp(status, serializedAllTrialInformation);
}
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java
index ea151cc6b12..d431e147d57 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java
@@ -131,6 +131,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListReq;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetLocationForTriggerResp;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPathsSetTemplatesResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipePluginTableResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetRegionIdReq;
@@ -161,8 +163,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp;
import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp;
import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp;
@@ -1833,11 +1835,11 @@ public class ConfigManager implements IManager {
}
@Override
- public TShowTrailResp showTrail(TShowTrailReq req) {
+ public TShowTrialResp showTrial(TShowTrialReq req) {
TSStatus status = confirmLeader();
return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
- ? modelManager.showTrail(req)
- : new TShowTrailResp(status, Collections.emptyList());
+ ? modelManager.showTrial(req)
+ : new TShowTrialResp(status, Collections.emptyList());
}
@Override
@@ -1856,6 +1858,14 @@ public class ConfigManager implements IManager {
: status;
}
+ @Override
+ public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) {
+ TSStatus status = confirmLeader();
+ return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
+ ? modelManager.getModelInfo(req)
+ : new TGetModelInfoResp(status);
+ }
+
@Override
public TSStatus setSpaceQuota(TSetSpaceQuotaReq req) {
TSStatus status = confirmLeader();
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/IManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/IManager.java
index 77c47b4cf7c..245fa73495c 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/IManager.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/IManager.java
@@ -76,6 +76,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListReq;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetLocationForTriggerResp;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPathsSetTemplatesResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipePluginTableResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetRegionIdReq;
@@ -104,8 +106,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp;
import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq;
import org.apache.iotdb.confignode.rpc.thrift.TUpdateModelInfoReq;
@@ -651,8 +653,8 @@ public interface IManager {
/** Return the model table. */
TShowModelResp showModel(TShowModelReq req);
- /** Return the trail table. */
- TShowTrailResp showTrail(TShowTrailReq req);
+ /** Return the trial table. */
+ TShowTrialResp showTrial(TShowTrialReq req);
/** Update the model info. */
TSStatus updateModelInfo(TUpdateModelInfoReq req);
@@ -660,6 +662,9 @@ public interface IManager {
/** Update the model state. */
TSStatus updateModelState(TUpdateModelStateReq req);
+ /** Update the model state */
+ TGetModelInfoResp getModelInfo(TGetModelInfoReq req);
+
/** Set space quota. */
TSStatus setSpaceQuota(TSetSpaceQuotaReq req);
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java
index 136d3849a88..288e1c2263c 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java
@@ -20,31 +20,49 @@
package org.apache.iotdb.confignode.manager;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.common.rpc.thrift.TaskType;
+import org.apache.iotdb.commons.model.ForecastModeInformation;
import org.apache.iotdb.commons.model.ModelInformation;
+import org.apache.iotdb.confignode.consensus.request.read.model.GetModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.read.model.ShowModelPlan;
-import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrailPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrialPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelStatePlan;
-import org.apache.iotdb.confignode.consensus.response.ModelTableResp;
-import org.apache.iotdb.confignode.consensus.response.TrailTableResp;
+import org.apache.iotdb.confignode.consensus.response.model.GetModelInfoResp;
+import org.apache.iotdb.confignode.consensus.response.model.ModelTableResp;
+import org.apache.iotdb.confignode.consensus.response.model.TrialTableResp;
import org.apache.iotdb.confignode.persistence.ModelInfo;
import org.apache.iotdb.confignode.rpc.thrift.TCreateModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropModelReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TUpdateModelInfoReq;
import org.apache.iotdb.confignode.rpc.thrift.TUpdateModelStateReq;
import org.apache.iotdb.consensus.common.DataSet;
import org.apache.iotdb.consensus.exception.ConsensusException;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.util.Arrays;
import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.apache.iotdb.commons.model.ForecastModeInformation.DEFAULT_INPUT_LENGTH;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.DEFAULT_PREDICT_LENGTH;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.INPUT_LENGTH;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.INPUT_TYPE_LIST;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.PREDICT_INDEX_LIST;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.PREDICT_LENGTH;
public class ModelManager {
@@ -63,15 +81,39 @@ public class ModelManager {
}
public TSStatus createModel(TCreateModelReq req) {
- ModelInformation modelInformation =
- new ModelInformation(
- req.getModelId(),
- req.getModelTask(),
- req.getModelType(),
- req.isIsAuto(),
- req.getQueryExpressions(),
- req.getQueryFilter());
- return configManager.getProcedureManager().createModel(modelInformation, req.getModelConfigs());
+ TaskType taskType = req.getTaskType();
+
+ Map<String, String> options = req.getOptions();
+ ModelInformation modelInformation;
+ if (taskType == TaskType.FORECAST) {
+ String inputTypeListStr = options.get(INPUT_TYPE_LIST);
+ List<TSDataType> inputTypeList =
+ Arrays.stream(inputTypeListStr.substring(1, inputTypeListStr.length() - 1).split(","))
+ .map(s -> TSDataType.valueOf(s.toUpperCase().trim()))
+ .collect(Collectors.toList());
+
+ String predictIndexListStr = options.get(PREDICT_INDEX_LIST);
+ List<Integer> predictIndexList =
+ Arrays.stream(
+ predictIndexListStr.substring(1, predictIndexListStr.length() - 1).split(","))
+ .map(s -> Integer.valueOf(s.trim()))
+ .collect(Collectors.toList());
+
+ modelInformation =
+ new ForecastModeInformation(
+ req.getModelId(),
+ req.getOptions(),
+ req.getDatasetFetchSQL().replace("\n", "").replace("\t", " "),
+ inputTypeList,
+ predictIndexList,
+ Integer.parseInt(options.getOrDefault(INPUT_LENGTH, DEFAULT_INPUT_LENGTH)),
+ Integer.parseInt(options.getOrDefault(PREDICT_LENGTH, DEFAULT_PREDICT_LENGTH)));
+ } else {
+ throw new IllegalArgumentException("Invalid task type: " + taskType);
+ }
+ return configManager
+ .getProcedureManager()
+ .createModel(modelInformation, req.getHyperparameters());
}
public TSStatus dropModel(TDropModelReq req) {
@@ -119,7 +161,7 @@ public class ModelManager {
res.setMessage(e.getMessage());
return new TShowModelResp(res, Collections.emptyList());
} catch (IOException e) {
- LOGGER.error("Fail to get ModelTable", e);
+ LOGGER.warn("Fail to get ModelTable", e);
return new TShowModelResp(
new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
.setMessage(e.getMessage()),
@@ -127,23 +169,36 @@ public class ModelManager {
}
}
- public TShowTrailResp showTrail(TShowTrailReq req) {
+ public TShowTrialResp showTrial(TShowTrialReq req) {
try {
- DataSet response = configManager.getConsensusManager().read(new ShowTrailPlan(req));
- return ((TrailTableResp) response).convertToThriftResponse();
+ DataSet response = configManager.getConsensusManager().read(new ShowTrialPlan(req));
+ return ((TrialTableResp) response).convertToThriftResponse();
} catch (ConsensusException e) {
LOGGER.warn(
- String.format("Unexpected error happened while showing trail %s: ", req.getModelId()), e);
+ String.format("Unexpected error happened while showing trial %s: ", req.getModelId()), e);
// consensus layer related errors
TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode());
res.setMessage(e.getMessage());
- return new TShowTrailResp(res, Collections.emptyList());
+ return new TShowTrialResp(res, Collections.emptyList());
} catch (IOException e) {
- LOGGER.error("Fail to get TrailTable", e);
- return new TShowTrailResp(
+ LOGGER.warn("Fail to get TrailTable", e);
+ return new TShowTrialResp(
new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
.setMessage(e.getMessage()),
Collections.emptyList());
}
}
+
+ public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) {
+ try {
+ DataSet response = configManager.getConsensusManager().read(new GetModelInfoPlan(req));
+ return ((GetModelInfoResp) response).convertToThriftResponse();
+ } catch (ConsensusException e) {
+ LOGGER.warn("Unexpected error happened while getting model: ", e);
+ // consensus layer related errors
+ TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode());
+ res.setMessage(e.getMessage());
+ return new TGetModelInfoResp(res);
+ }
+ }
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java
index c4b504abdb2..e3496fadf0f 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java
@@ -704,9 +704,10 @@ public class ProcedureManager {
return statusList.get(0);
}
- public TSStatus createModel(ModelInformation modelInformation, Map<String, String> modelConfigs) {
+ public TSStatus createModel(
+ ModelInformation modelInformation, Map<String, String> hyperparameters) {
long procedureId =
- executor.submitProcedure(new CreateModelProcedure(modelInformation, modelConfigs));
+ executor.submitProcedure(new CreateModelProcedure(modelInformation, hyperparameters));
List<TSStatus> statusList = new ArrayList<>();
boolean isSucceed =
waitingProcedureFinished(Collections.singletonList(procedureId), statusList);
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/ModelInfo.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/ModelInfo.java
index 3c72e09570f..89996da9717 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/ModelInfo.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/ModelInfo.java
@@ -22,17 +22,20 @@ package org.apache.iotdb.confignode.persistence;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.commons.model.ModelTable;
-import org.apache.iotdb.commons.model.TrailInformation;
+import org.apache.iotdb.commons.model.TrialInformation;
import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
+import org.apache.iotdb.confignode.consensus.request.read.model.GetModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.read.model.ShowModelPlan;
-import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrailPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrialPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.CreateModelPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.DropModelPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelStatePlan;
-import org.apache.iotdb.confignode.consensus.response.ModelTableResp;
-import org.apache.iotdb.confignode.consensus.response.TrailTableResp;
+import org.apache.iotdb.confignode.consensus.response.model.GetModelInfoResp;
+import org.apache.iotdb.confignode.consensus.response.model.ModelTableResp;
+import org.apache.iotdb.confignode.consensus.response.model.TrialTableResp;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.tsfile.utils.PublicBAOS;
import org.apache.thrift.TException;
import org.slf4j.Logger;
@@ -40,11 +43,14 @@ import org.slf4j.LoggerFactory;
import javax.annotation.concurrent.ThreadSafe;
+import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
-import java.util.concurrent.locks.ReentrantLock;
+import java.nio.ByteBuffer;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
@ThreadSafe
public class ModelInfo implements SnapshotProcessor {
@@ -55,48 +61,76 @@ public class ModelInfo implements SnapshotProcessor {
private ModelTable modelTable;
- private final ReentrantLock modelTableLock = new ReentrantLock();
+ private final ReadWriteLock modelTableLock = new ReentrantReadWriteLock();
public ModelInfo() {
this.modelTable = new ModelTable();
}
- public void acquireModelTableLock() {
- LOGGER.info("acquire ModelTableLock");
- modelTableLock.lock();
+ public void acquireModelTableReadLock() {
+ LOGGER.info("acquire ModelTableReadLock");
+ modelTableLock.readLock().lock();
}
- public void releaseModelTableLock() {
- LOGGER.info("release ModelTableLock");
- modelTableLock.unlock();
+ public void releaseModelTableReadLock() {
+ LOGGER.info("release ModelTableReadLock");
+ modelTableLock.readLock().unlock();
+ }
+
+ public void acquireModelTableWriteLock() {
+ LOGGER.info("acquire ModelTableWriteLock");
+ modelTableLock.writeLock().lock();
+ }
+
+ public void releaseModelTableWriteLock() {
+ LOGGER.info("release ModelTableWriteLock");
+ modelTableLock.writeLock().unlock();
}
public TSStatus createModel(CreateModelPlan plan) {
try {
+ acquireModelTableWriteLock();
ModelInformation modelInformation = plan.getModelInformation();
- modelTable.addModel(modelInformation);
+ if (modelTable.containsModel(modelInformation.getModelId())) {
+ return new TSStatus(TSStatusCode.MODEL_EXIST_ERROR.getStatusCode())
+ .setMessage(
+ String.format(
+ "model [%s] has already been created.", modelInformation.getModelId()));
+ } else {
+ modelTable.addModel(modelInformation);
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+
} catch (Exception e) {
final String errorMessage =
String.format(
"Failed to add model [%s] in ModelTable on Config Nodes, because of %s",
plan.getModelInformation().getModelId(), e);
LOGGER.warn(errorMessage, e);
- return new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
- .setMessage(errorMessage);
+ return new TSStatus(TSStatusCode.CREATE_MODEL_ERROR.getStatusCode()).setMessage(errorMessage);
+ } finally {
+ releaseModelTableWriteLock();
}
- return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
}
public TSStatus dropModel(DropModelPlan plan) {
+ acquireModelTableWriteLock();
String modelId = plan.getModelId();
+ TSStatus status;
if (modelTable.containsModel(modelId)) {
modelTable.removeModel(modelId);
+ status = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ } else {
+ status =
+ new TSStatus(TSStatusCode.DROP_MODEL_ERROR.getStatusCode())
+ .setMessage(String.format("model [%s] has not been created.", modelId));
}
- return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ releaseModelTableWriteLock();
+ return status;
}
public ModelTableResp showModel(ShowModelPlan plan) {
- acquireModelTableLock();
+ acquireModelTableReadLock();
try {
ModelTableResp modelTableResp =
new ModelTableResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
@@ -115,17 +149,17 @@ public class ModelInfo implements SnapshotProcessor {
new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
.setMessage(e.getMessage()));
} finally {
- releaseModelTableLock();
+ releaseModelTableReadLock();
}
}
- public TrailTableResp showTrail(ShowTrailPlan plan) {
- acquireModelTableLock();
+ public TrialTableResp showTrial(ShowTrialPlan plan) {
+ acquireModelTableReadLock();
try {
String modelId = plan.getModelId();
ModelInformation modelInformation = modelTable.getModelInformationById(modelId);
if (modelInformation == null) {
- return new TrailTableResp(
+ return new TrialTableResp(
new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
.setMessage(
String.format(
@@ -133,51 +167,74 @@ public class ModelInfo implements SnapshotProcessor {
modelId)));
}
- TrailTableResp trailTableResp =
- new TrailTableResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
- if (plan.isSetTrailId()) {
- TrailInformation trailInformation =
- modelInformation.getTrailInformationById(plan.getTrailId());
- if (trailInformation != null) {
- trailTableResp.addTrailInformation(trailInformation);
+ TrialTableResp trialTableResp =
+ new TrialTableResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+ if (plan.isSetTrialId()) {
+ TrialInformation trialInformation =
+ modelInformation.getTrialInformationById(plan.getTrialId());
+ if (trialInformation != null) {
+ trialTableResp.addTrialInformation(trialInformation);
}
} else {
- trailTableResp.addTrailInformation(modelInformation.getAllTrailInformation());
+ trialTableResp.addTrialInformation(modelInformation.getAllTrialInformation());
}
- return trailTableResp;
+ return trialTableResp;
} catch (IOException e) {
LOGGER.warn("Fail to get TrailTable", e);
- return new TrailTableResp(
+ return new TrialTableResp(
new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
.setMessage(e.getMessage()));
} finally {
- releaseModelTableLock();
+ releaseModelTableReadLock();
+ }
+ }
+
+ public GetModelInfoResp getModelInfo(GetModelInfoPlan plan) {
+ acquireModelTableReadLock();
+ try {
+ GetModelInfoResp getModelInfoResp =
+ new GetModelInfoResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+ ModelInformation modelInformation = modelTable.getModelInformationById(plan.getModelId());
+ if (modelInformation != null) {
+ PublicBAOS buffer = new PublicBAOS();
+ DataOutputStream stream = new DataOutputStream(buffer);
+ modelInformation.serialize(stream);
+ getModelInfoResp.setModelInfo(ByteBuffer.wrap(buffer.getBuf(), 0, buffer.size()));
+ }
+ return getModelInfoResp;
+ } catch (IOException e) {
+ LOGGER.warn("Fail to get model info", e);
+ return new GetModelInfoResp(
+ new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
+ .setMessage(e.getMessage()));
+ } finally {
+ releaseModelTableReadLock();
}
}
public TSStatus updateModelInfo(UpdateModelInfoPlan plan) {
- acquireModelTableLock();
+ acquireModelTableWriteLock();
try {
String modelId = plan.getModelId();
if (modelTable.containsModel(modelId)) {
- modelTable.updateModel(modelId, plan.getTrailId(), plan.getModelInfo());
+ modelTable.updateModel(modelId, plan.getTrialId(), plan.getModelInfo());
}
return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
} finally {
- releaseModelTableLock();
+ releaseModelTableWriteLock();
}
}
public TSStatus updateModelState(UpdateModelStatePlan plan) {
- acquireModelTableLock();
+ acquireModelTableWriteLock();
try {
String modelId = plan.getModelId();
if (modelTable.containsModel(modelId)) {
- modelTable.updateState(modelId, plan.getState(), plan.getBestTrailId());
+ modelTable.updateState(modelId, plan.getState(), plan.getBestTrialId());
}
return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
} finally {
- releaseModelTableLock();
+ releaseModelTableWriteLock();
}
}
@@ -191,12 +248,12 @@ public class ModelInfo implements SnapshotProcessor {
return false;
}
- acquireModelTableLock();
+ acquireModelTableReadLock();
try (FileOutputStream fileOutputStream = new FileOutputStream(snapshotFile)) {
modelTable.serialize(fileOutputStream);
return true;
} finally {
- releaseModelTableLock();
+ releaseModelTableReadLock();
}
}
@@ -209,16 +266,12 @@ public class ModelInfo implements SnapshotProcessor {
snapshotFile.getAbsolutePath());
return;
}
- acquireModelTableLock();
+ acquireModelTableWriteLock();
try (FileInputStream fileInputStream = new FileInputStream(snapshotFile)) {
modelTable.clear();
modelTable = ModelTable.deserialize(fileInputStream);
} finally {
- releaseModelTableLock();
+ releaseModelTableWriteLock();
}
}
-
- public boolean isModelExist(String modelId) {
- return modelTable.containsModel(modelId);
- }
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java
index 2aff2dc81f0..c859c7ef530 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java
@@ -31,8 +31,9 @@ import org.apache.iotdb.confignode.consensus.request.read.database.CountDatabase
import org.apache.iotdb.confignode.consensus.request.read.database.GetDatabasePlan;
import org.apache.iotdb.confignode.consensus.request.read.datanode.GetDataNodeConfigurationPlan;
import org.apache.iotdb.confignode.consensus.request.read.function.GetUDFJarPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.GetModelInfoPlan;
import org.apache.iotdb.confignode.consensus.request.read.model.ShowModelPlan;
-import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrailPlan;
+import org.apache.iotdb.confignode.consensus.request.read.model.ShowTrialPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.CountTimeSlotListPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.GetDataPartitionPlan;
import org.apache.iotdb.confignode.consensus.request.read.partition.GetNodePathsPartitionPlan;
@@ -279,8 +280,10 @@ public class ConfigPlanExecutor {
return udfInfo.getUDFJar((GetUDFJarPlan) req);
case ShowModel:
return modelInfo.showModel((ShowModelPlan) req);
- case ShowTrail:
- return modelInfo.showTrail((ShowTrailPlan) req);
+ case ShowTrial:
+ return modelInfo.showTrial((ShowTrialPlan) req);
+ case GetModelInfo:
+ return modelInfo.getModelInfo((GetModelInfoPlan) req);
case GetPipePluginTable:
return pipeInfo.getPipePluginInfo().showPipePlugins();
case GetPipePluginJar:
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/CreateModelProcedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/CreateModelProcedure.java
index 2891fe80049..d118f7382f6 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/CreateModelProcedure.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/CreateModelProcedure.java
@@ -20,19 +20,20 @@
package org.apache.iotdb.confignode.procedure.impl.model;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClient;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClientManager;
+import org.apache.iotdb.commons.client.mlnode.MLNodeInfo;
import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.commons.model.exception.ModelManagementException;
import org.apache.iotdb.confignode.consensus.request.write.model.CreateModelPlan;
import org.apache.iotdb.confignode.consensus.request.write.model.DropModelPlan;
import org.apache.iotdb.confignode.manager.ConfigManager;
-import org.apache.iotdb.confignode.persistence.ModelInfo;
import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
import org.apache.iotdb.confignode.procedure.impl.node.AbstractNodeProcedure;
import org.apache.iotdb.confignode.procedure.state.model.CreateModelState;
import org.apache.iotdb.confignode.procedure.store.ProcedureType;
import org.apache.iotdb.consensus.exception.ConsensusException;
-import org.apache.iotdb.db.protocol.client.MLNodeClient;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
@@ -52,16 +53,17 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
private static final int RETRY_THRESHOLD = 1;
private ModelInformation modelInformation;
- private Map<String, String> modelConfigs;
+ private Map<String, String> hyperparameters;
public CreateModelProcedure() {
super();
}
- public CreateModelProcedure(ModelInformation modelInformation, Map<String, String> modelConfigs) {
+ public CreateModelProcedure(
+ ModelInformation modelInformation, Map<String, String> hyperparameters) {
super();
this.modelInformation = modelInformation;
- this.modelConfigs = modelConfigs;
+ this.hyperparameters = hyperparameters;
}
@Override
@@ -72,58 +74,14 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
try {
switch (state) {
case INIT:
- LOGGER.info("Start to create model [{}]", modelInformation.getModelId());
-
- ModelInfo modelInfo = env.getConfigManager().getModelManager().getModelInfo();
- modelInfo.acquireModelTableLock();
- String modelId = modelInformation.getModelId();
- if (modelInfo.isModelExist(modelId)) {
- throw new ModelManagementException(
- String.format(
- "Failed to create model [%s], the same name model has been created", modelId));
- }
- setNextState(CreateModelState.VALIDATED);
+ createModel(env);
break;
-
- case VALIDATED:
- ConfigManager configManager = env.getConfigManager();
- modelId = modelInformation.getModelId();
-
- LOGGER.info("Start to add model [{}] in ModelTable on Config Nodes", modelId);
-
- TSStatus response =
- configManager.getConsensusManager().write(new CreateModelPlan(modelInformation));
- if (response.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- throw new ModelManagementException(
- String.format(
- "Failed to add model [%s] in ModelTable on Config Nodes: %s",
- modelId, response.getMessage()));
- }
-
- setNextState(CreateModelState.CONFIG_NODE_ACTIVE);
- break;
-
- case CONFIG_NODE_ACTIVE:
- LOGGER.info("Start to train model [{}] on ML Node", modelInformation.getModelId());
-
- try (MLNodeClient client = new MLNodeClient()) {
- TSStatus status = client.createTrainingTask(modelInformation, modelConfigs);
- if (status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- throw new TException(status.getMessage());
- }
-
- setNextState(CreateModelState.ML_NODE_ACTIVE);
- } catch (TException e) {
- throw new ModelManagementException(
- String.format(
- "Fail to start training model [%s] on ML Node: %s",
- modelInformation.getModelId(), e.getMessage()));
- }
- break;
-
case ML_NODE_ACTIVE:
- env.getConfigManager().getModelManager().getModelInfo().releaseModelTableLock();
+ trainModel();
return Flow.NO_MORE_STATE;
+ default:
+ throw new UnsupportedOperationException(
+ String.format("Unknown state during executing createModelProcedure, %s", state));
}
} catch (Exception e) {
if (isRollbackSupported(state)) {
@@ -147,19 +105,47 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
return Flow.HAS_MORE_STATE;
}
+ private void createModel(ConfigNodeProcedureEnv env) throws ConsensusException {
+ LOGGER.info("Start to create model [{}]", modelInformation.getModelId());
+
+ ConfigManager configManager = env.getConfigManager();
+ TSStatus response =
+ configManager.getConsensusManager().write(new CreateModelPlan(modelInformation));
+ if (response.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ throw new ModelManagementException(
+ String.format(
+ "Failed to add model [%s] in ModelTable on Config Nodes: %s",
+ modelInformation.getModelId(), response.getMessage()));
+ }
+ setNextState(CreateModelState.ML_NODE_ACTIVE);
+ }
+
+ private void trainModel() {
+ LOGGER.info("Start to train model [{}] on ML Node", modelInformation.getModelId());
+
+ try (MLNodeClient client =
+ MLNodeClientManager.getInstance().borrowClient(MLNodeInfo.endPoint)) {
+ TSStatus status = client.createTrainingTask(modelInformation, hyperparameters);
+ if (status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ throw new TException(status.getMessage());
+ }
+ } catch (Exception e) {
+ throw new ModelManagementException(
+ String.format(
+ "Fail to start training model [%s] on ML Node: %s",
+ modelInformation.getModelId(), e.getMessage()));
+ }
+ }
+
@Override
protected void rollbackState(ConfigNodeProcedureEnv env, CreateModelState state)
throws IOException, InterruptedException, ProcedureException {
switch (state) {
case INIT:
- LOGGER.info("Start [INIT] rollback of model [{}]", modelInformation.getModelId());
-
- env.getConfigManager().getModelManager().getModelInfo().releaseModelTableLock();
+ // do nothing
break;
-
- case VALIDATED:
+ case ML_NODE_ACTIVE:
LOGGER.info("Start [VALIDATED] rollback of model [{}]", modelInformation.getModelId());
-
try {
env.getConfigManager()
.getConsensusManager()
@@ -168,10 +154,19 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
LOGGER.warn("Failed in the write API executing the consensus layer due to: ", e);
}
break;
-
default:
break;
}
+ if (Objects.requireNonNull(state) == CreateModelState.INIT) {
+ LOGGER.info("Start [INIT] rollback of model [{}]", modelInformation.getModelId());
+ try {
+ env.getConfigManager()
+ .getConsensusManager()
+ .write(new DropModelPlan(modelInformation.getModelId()));
+ } catch (ConsensusException e) {
+ LOGGER.warn("Failed in the write API executing the consensus layer due to: ", e);
+ }
+ }
}
@Override
@@ -199,14 +194,14 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
stream.writeShort(ProcedureType.CREATE_MODEL_PROCEDURE.getTypeCode());
super.serialize(stream);
modelInformation.serialize(stream);
- ReadWriteIOUtils.write(modelConfigs, stream);
+ ReadWriteIOUtils.write(hyperparameters, stream);
}
@Override
public void deserialize(ByteBuffer byteBuffer) {
super.deserialize(byteBuffer);
modelInformation = ModelInformation.deserialize(byteBuffer);
- modelConfigs = ReadWriteIOUtils.readMap(byteBuffer);
+ hyperparameters = ReadWriteIOUtils.readMap(byteBuffer);
}
@Override
@@ -216,13 +211,13 @@ public class CreateModelProcedure extends AbstractNodeProcedure<CreateModelState
return thatProc.getProcId() == this.getProcId()
&& thatProc.getState() == this.getState()
&& thatProc.modelInformation.equals(this.modelInformation)
- && thatProc.modelConfigs.equals(this.modelConfigs);
+ && thatProc.hyperparameters.equals(this.hyperparameters);
}
return false;
}
@Override
public int hashCode() {
- return Objects.hash(getProcId(), getState(), modelInformation, modelConfigs);
+ return Objects.hash(getProcId(), getState(), modelInformation, hyperparameters);
}
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/DropModelProcedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/DropModelProcedure.java
index 6c1f92c2d4c..0cf7ea2dae3 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/DropModelProcedure.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/DropModelProcedure.java
@@ -21,17 +21,18 @@ package org.apache.iotdb.confignode.procedure.impl.model;
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClient;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClientManager;
+import org.apache.iotdb.commons.client.mlnode.MLNodeInfo;
import org.apache.iotdb.commons.model.exception.ModelManagementException;
import org.apache.iotdb.confignode.client.DataNodeRequestType;
import org.apache.iotdb.confignode.client.sync.SyncDataNodeClientPool;
import org.apache.iotdb.confignode.consensus.request.write.model.DropModelPlan;
-import org.apache.iotdb.confignode.persistence.ModelInfo;
import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
import org.apache.iotdb.confignode.procedure.impl.node.AbstractNodeProcedure;
import org.apache.iotdb.confignode.procedure.state.model.DropModelState;
import org.apache.iotdb.confignode.procedure.store.ProcedureType;
-import org.apache.iotdb.db.protocol.client.MLNodeClient;
import org.apache.iotdb.mpp.rpc.thrift.TDeleteModelMetricsReq;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
@@ -70,66 +71,7 @@ public class DropModelProcedure extends AbstractNodeProcedure<DropModelState> {
try {
switch (state) {
case INIT:
- LOGGER.info("Start to drop model [{}]", modelId);
-
- ModelInfo modelInfo = env.getConfigManager().getModelManager().getModelInfo();
- modelInfo.acquireModelTableLock();
- if (!modelInfo.isModelExist(modelId)) {
- throw new ModelManagementException(
- String.format(
- "Failed to drop model [%s], this model has not been created", modelId));
- }
- setNextState(DropModelState.VALIDATED);
- break;
-
- case VALIDATED:
- LOGGER.info("Start to drop model metrics [{}] on Data Nodes", modelId);
-
- Optional<TDataNodeLocation> targetDataNode =
- env.getConfigManager().getNodeManager().getLowestLoadDataNode();
- if (!targetDataNode.isPresent()) {
- // no usable DataNode
- throw new ModelManagementException(
- String.format("Failed to drop model [%s], there is no RUNNING DataNode", modelId));
- }
-
- TSStatus status =
- SyncDataNodeClientPool.getInstance()
- .sendSyncRequestToDataNodeWithRetry(
- targetDataNode.get().getInternalEndPoint(),
- new TDeleteModelMetricsReq(modelId),
- DataNodeRequestType.DELETE_MODEL_METRICS);
- if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- throw new ModelManagementException(
- String.format(
- "Failed to drop model [%s], fail to delete metrics: %s",
- modelId, status.getMessage()));
- }
-
- setNextState(DropModelState.DATA_NODE_DROPPED);
- break;
-
- case DATA_NODE_DROPPED:
- LOGGER.info("Start to drop model file [{}] on Ml Node", modelId);
-
- try (MLNodeClient client = new MLNodeClient()) {
- status = client.deleteModel(modelId);
- if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- throw new TException(status.getMessage());
- }
- } catch (TException e) {
- throw new ModelManagementException(
- String.format(
- "Failed to drop model [%s], fail to delete model on MLNode: %s",
- modelId, e.getMessage()));
- }
-
- setNextState(DropModelState.ML_NODE_DROPPED);
- break;
-
- case ML_NODE_DROPPED:
LOGGER.info("Start to drop model [{}] on Config Nodes", modelId);
-
TSStatus response =
env.getConfigManager().getConsensusManager().write(new DropModelPlan(modelId));
if (response.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
@@ -138,13 +80,17 @@ public class DropModelProcedure extends AbstractNodeProcedure<DropModelState> {
"Failed to drop model [%s], fail to drop model on Config Nodes: %s",
modelId, response.getMessage()));
}
-
- setNextState(DropModelState.CONFIG_NODE_DROPPED);
+ setNextState(DropModelState.DATA_NODE_DROPPED);
break;
-
- case CONFIG_NODE_DROPPED:
- env.getConfigManager().getModelManager().getModelInfo().releaseModelTableLock();
+ case DATA_NODE_DROPPED:
+ deleteMetricsOnDataNode(env);
+ break;
+ case ML_NODE_DROPPED:
+ dropModelOnMLNode();
return Flow.NO_MORE_STATE;
+ default:
+ throw new UnsupportedOperationException(
+ String.format("Unknown state during executing dropModelProcedure, %s", state));
}
} catch (Exception e) {
if (isRollbackSupported(state)) {
@@ -164,14 +110,57 @@ public class DropModelProcedure extends AbstractNodeProcedure<DropModelState> {
return Flow.HAS_MORE_STATE;
}
+ private void dropModelOnMLNode() {
+ LOGGER.info("Start to drop model file [{}] on Ml Node", modelId);
+ try (MLNodeClient client =
+ MLNodeClientManager.getInstance().borrowClient(MLNodeInfo.endPoint)) {
+ TSStatus status = client.deleteModel(modelId);
+ if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ throw new TException(status.getMessage());
+ }
+ } catch (Exception e) {
+ throw new ModelManagementException(
+ String.format(
+ "Failed to drop model [%s], fail to delete model on MLNode: %s",
+ modelId, e.getMessage()));
+ }
+ }
+
+ private void deleteMetricsOnDataNode(ConfigNodeProcedureEnv env) {
+ LOGGER.info("Start to drop model metrics [{}] on Data Nodes", modelId);
+
+ Optional<TDataNodeLocation> targetDataNode =
+ env.getConfigManager().getNodeManager().getLowestLoadDataNode();
+ if (!targetDataNode.isPresent()) {
+ // no usable DataNode
+ throw new ModelManagementException(
+ String.format("Failed to drop model [%s], there is no RUNNING DataNode", modelId));
+ }
+
+ TSStatus status =
+ SyncDataNodeClientPool.getInstance()
+ .sendSyncRequestToDataNodeWithRetry(
+ targetDataNode.get().getInternalEndPoint(),
+ new TDeleteModelMetricsReq(modelId),
+ DataNodeRequestType.DELETE_MODEL_METRICS);
+ if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ // It's ok for the next state if the metric path doesn't exit
+ if (status.getCode() != TSStatusCode.PATH_NOT_EXIST.getStatusCode()) {
+ throw new ModelManagementException(
+ String.format(
+ "Failed to drop model [%s], fail to delete metrics: %s",
+ modelId, status.getMessage()));
+ }
+ LOGGER.warn("The path of metrics does not exist on model {}", modelId);
+ }
+
+ setNextState(DropModelState.ML_NODE_DROPPED);
+ }
+
@Override
protected void rollbackState(ConfigNodeProcedureEnv env, DropModelState state)
throws IOException, InterruptedException, ProcedureException {
- if (state == DropModelState.INIT) {
- LOGGER.info("Start [INIT] rollback of model [{}]", modelId);
-
- env.getConfigManager().getModelManager().getModelInfo().releaseModelTableLock();
- }
+ // no need to rollback
}
@Override
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
index 304d209817f..aad03181187 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
@@ -21,7 +21,5 @@ package org.apache.iotdb.confignode.procedure.state.model;
public enum CreateModelState {
INIT,
- VALIDATED,
- CONFIG_NODE_ACTIVE,
ML_NODE_ACTIVE
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
index 01131f6f125..2c1430cfa2e 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
@@ -112,6 +112,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListReq;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetLocationForTriggerResp;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPathsSetTemplatesResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipePluginTableResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipeSinkReq;
@@ -152,8 +154,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp;
import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp;
import org.apache.iotdb.confignode.rpc.thrift.TSystemConfigurationResp;
@@ -981,8 +983,8 @@ public class ConfigNodeRPCServiceProcessor implements IConfigNodeRPCService.Ifac
}
@Override
- public TShowTrailResp showTrail(TShowTrailReq req) {
- return configManager.showTrail(req);
+ public TShowTrialResp showTrial(TShowTrialReq req) {
+ return configManager.showTrial(req);
}
@Override
@@ -995,6 +997,11 @@ public class ConfigNodeRPCServiceProcessor implements IConfigNodeRPCService.Ifac
return configManager.updateModelState(req);
}
+ @Override
+ public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) {
+ return configManager.getModelInfo(req);
+ }
+
@Override
public TSStatus setSpaceQuota(TSetSpaceQuotaReq req) throws TException {
return configManager.setSpaceQuota(req);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java
index 543256d421f..890b0a46c76 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java
@@ -80,6 +80,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListReq;
import org.apache.iotdb.confignode.rpc.thrift.TGetJarInListResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetLocationForTriggerResp;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPathsSetTemplatesResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipePluginTableResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipeSinkReq;
@@ -120,8 +122,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp;
import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp;
import org.apache.iotdb.confignode.rpc.thrift.TSystemConfigurationResp;
@@ -954,9 +956,9 @@ public class ConfigNodeClient implements IConfigNodeRPCService.Iface, ThriftClie
}
@Override
- public TShowTrailResp showTrail(TShowTrailReq req) throws TException {
+ public TShowTrialResp showTrial(TShowTrialReq req) throws TException {
return executeRemoteCallWithRetry(
- () -> client.showTrail(req), resp -> !updateConfigNodeLeader(resp.status));
+ () -> client.showTrial(req), resp -> !updateConfigNodeLeader(resp.status));
}
@Override
@@ -969,6 +971,22 @@ public class ConfigNodeClient implements IConfigNodeRPCService.Iface, ThriftClie
throw new TException(new UnsupportedOperationException().getCause());
}
+ @Override
+ public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) throws TException {
+ for (int i = 0; i < RETRY_NUM; i++) {
+ try {
+ TGetModelInfoResp resp = client.getModelInfo(req);
+ if (!updateConfigNodeLeader(resp.getStatus())) {
+ return resp;
+ }
+ } catch (TException e) {
+ configLeader = null;
+ }
+ waitAndReconnect();
+ }
+ throw new TException(MSG_RECONNECTION_FAIL);
+ }
+
@Override
public TSStatus setSpaceQuota(TSetSpaceQuotaReq req) throws TException {
return executeRemoteCallWithRetry(
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/MLNodeRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/MLNodeRPCServiceImpl.java
index d881124f508..46f1a966112 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/MLNodeRPCServiceImpl.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/MLNodeRPCServiceImpl.java
@@ -34,8 +34,8 @@ import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaFetcher;
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
import org.apache.iotdb.db.queryengine.plan.parser.StatementGenerator;
+import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
-import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement;
import org.apache.iotdb.db.utils.ErrorHandlingUtils;
import org.apache.iotdb.db.utils.QueryDataSetUtils;
import org.apache.iotdb.db.utils.SetThreadName;
@@ -46,6 +46,7 @@ import org.apache.iotdb.mpp.rpc.thrift.TFetchTimeseriesResp;
import org.apache.iotdb.mpp.rpc.thrift.TFetchWindowBatchReq;
import org.apache.iotdb.mpp.rpc.thrift.TFetchWindowBatchResp;
import org.apache.iotdb.mpp.rpc.thrift.TRecordModelMetricsReq;
+import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.tsfile.utils.Pair;
@@ -85,8 +86,15 @@ public class MLNodeRPCServiceImpl implements IMLNodeRPCServiceWithHandler {
TFetchTimeseriesResp resp = new TFetchTimeseriesResp();
Throwable t = null;
try {
- QueryStatement s =
- (QueryStatement) StatementGenerator.createStatement(req, session.getZoneId());
+
+ Statement s = StatementGenerator.createStatement(req, session.getZoneId());
+
+ if (s == null) {
+ resp.setStatus(
+ RpcUtils.getStatus(
+ TSStatusCode.SQL_PARSE_ERROR, "This operation type is not supported"));
+ return resp;
+ }
long queryId =
SESSION_MANAGER.requestQueryId(session, SESSION_MANAGER.requestStatementId(session));
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/ColumnHeaderConstant.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/ColumnHeaderConstant.java
index 5cd5e9a2e22..93a1635f3a5 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/ColumnHeaderConstant.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/ColumnHeaderConstant.java
@@ -179,9 +179,9 @@ public class ColumnHeaderConstant {
public static final String USER = "User";
public static final String READ_WRITE = "Read/Write";
- // column names for show models/trails
+ // column names for show models/trials
public static final String MODEL_ID = "ModelId";
- public static final String TRAIL_ID = "TrailId";
+ public static final String TRIAL_ID = "TrialId";
public static final String MODEL_TASK = "ModelTask";
public static final String MODEL_TYPE = "ModelType";
public static final String QUERY_BODY = "QueryBody";
@@ -459,12 +459,11 @@ public class ColumnHeaderConstant {
new ColumnHeader(MODEL_TYPE, TSDataType.TEXT),
new ColumnHeader(QUERY_BODY, TSDataType.TEXT),
new ColumnHeader(STATE, TSDataType.TEXT),
- new ColumnHeader(MODEL_PATH, TSDataType.TEXT),
new ColumnHeader(HYPERPARAMETER, TSDataType.TEXT));
- public static final List<ColumnHeader> showTrailsColumnHeaders =
+ public static final List<ColumnHeader> showTrialsColumnHeaders =
ImmutableList.of(
- new ColumnHeader(TRAIL_ID, TSDataType.TEXT),
+ new ColumnHeader(TRIAL_ID, TSDataType.TEXT),
new ColumnHeader(MODEL_PATH, TSDataType.TEXT),
new ColumnHeader(HYPERPARAMETER, TSDataType.TEXT));
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java
index 99b2eb33344..423626d227f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java
@@ -189,8 +189,8 @@ public class DatasetHeaderFactory {
return new DatasetHeader(ColumnHeaderConstant.showModelsColumnHeaders, true);
}
- public static DatasetHeader getShowTrailsHeader() {
- return new DatasetHeader(ColumnHeaderConstant.showTrailsColumnHeaders, true);
+ public static DatasetHeader getShowTrialsHeader() {
+ return new DatasetHeader(ColumnHeaderConstant.showTrialsColumnHeaders, true);
}
public static DatasetHeader getShowLogicalViewHeader() {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
index 5e9bd9dd7fa..97d6be05df8 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.db.queryengine.execution.fragment;
import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
import org.apache.iotdb.commons.concurrent.ThreadName;
import org.apache.iotdb.commons.concurrent.threadpool.ScheduledExecutorUtil;
+import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
import org.apache.iotdb.db.queryengine.common.QueryId;
@@ -78,6 +79,7 @@ public class FragmentInstanceManager {
private final CounterStat failedInstances = new CounterStat();
private final ExecutorService intoOperationExecutor;
+ private final ExecutorService modelInferenceExecutor;
private final MPPDataExchangeManager exchangeManager =
MPPDataExchangeService.getInstance().getMPPDataExchangeManager();
@@ -115,6 +117,11 @@ public class FragmentInstanceManager {
IoTDBThreadPoolFactory.newFixedThreadPool(
IoTDBDescriptor.getInstance().getConfig().getIntoOperationExecutionThreadCount(),
"into-operation-executor");
+
+ this.modelInferenceExecutor =
+ IoTDBThreadPoolFactory.newFixedThreadPool(
+ CommonDescriptor.getInstance().getConfig().getModelInferenceExecutionThreadCount(),
+ "model-inference-executor");
}
public int getInstanceContextSize() {
@@ -355,6 +362,10 @@ public class FragmentInstanceManager {
return intoOperationExecutor;
}
+ public ExecutorService getModelInferenceExecutor() {
+ return modelInferenceExecutor;
+ }
+
private static class InstanceHolder {
private InstanceHolder() {}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationUtil.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationUtil.java
index 0f4eb86776c..df773c88996 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationUtil.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationUtil.java
@@ -223,7 +223,7 @@ public class AggregationUtil {
return timeValueColumnsSizePerLine;
}
- private static long getOutputColumnSizePerLine(TSDataType tsDataType) {
+ public static long getOutputColumnSizePerLine(TSDataType tsDataType) {
switch (tsDataType) {
case INT32:
return IntColumn.SIZE_IN_BYTES_PER_POSITION;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/ml/ForecastOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/ml/ForecastOperator.java
new file mode 100644
index 00000000000..0753d08683c
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/ml/ForecastOperator.java
@@ -0,0 +1,241 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.execution.operator.process.ml;
+
+import org.apache.iotdb.commons.client.mlnode.MLNodeClient;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClientManager;
+import org.apache.iotdb.commons.client.mlnode.MLNodeInfo;
+import org.apache.iotdb.db.exception.ModelInferenceProcessException;
+import org.apache.iotdb.db.queryengine.execution.operator.Operator;
+import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
+import org.apache.iotdb.db.queryengine.execution.operator.process.ProcessOperator;
+import org.apache.iotdb.db.utils.TimestampPrecisionUtils;
+import org.apache.iotdb.mlnode.rpc.thrift.TForecastResp;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
+import org.apache.iotdb.tsfile.read.common.block.TsBlock;
+import org.apache.iotdb.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.iotdb.tsfile.read.common.block.column.ColumnBuilder;
+import org.apache.iotdb.tsfile.read.common.block.column.TimeColumnBuilder;
+import org.apache.iotdb.tsfile.read.common.block.column.TsBlockSerde;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static com.google.common.util.concurrent.Futures.successfulAsList;
+
+public class ForecastOperator implements ProcessOperator {
+
+ private final OperatorContext operatorContext;
+ private final Operator child;
+
+ private final String modelPath;
+ private final List<TSDataType> inputTypeList;
+ private final List<String> inputColumnNameList;
+ private final int expectedPredictLength;
+
+ private final TsBlockBuilder inputTsBlockBuilder;
+
+ private final ExecutorService modelInferenceExecutor;
+ private ListenableFuture<TForecastResp> forecastExecutionFuture;
+
+ private boolean finished = false;
+
+ private final long maxRetainedSize;
+ private final long maxReturnSize;
+
+ public ForecastOperator(
+ OperatorContext operatorContext,
+ Operator child,
+ String modelPath,
+ List<TSDataType> inputTypeList,
+ List<String> inputColumnNameList,
+ int expectedPredictLength,
+ ExecutorService modelInferenceExecutor,
+ long maxRetainedSize,
+ long maxReturnSize) {
+ this.operatorContext = operatorContext;
+ this.child = child;
+ this.modelPath = modelPath;
+ this.inputTypeList = inputTypeList;
+ this.inputColumnNameList = inputColumnNameList;
+ this.expectedPredictLength = expectedPredictLength;
+ this.inputTsBlockBuilder = new TsBlockBuilder(inputTypeList);
+ this.modelInferenceExecutor = modelInferenceExecutor;
+ this.maxRetainedSize = maxRetainedSize;
+ this.maxReturnSize = maxReturnSize;
+ }
+
+ @Override
+ public OperatorContext getOperatorContext() {
+ return operatorContext;
+ }
+
+ @Override
+ public ListenableFuture<?> isBlocked() {
+ ListenableFuture<?> childBlocked = child.isBlocked();
+ boolean executionDone = forecastExecutionDone();
+ if (executionDone && childBlocked.isDone()) {
+ return NOT_BLOCKED;
+ } else if (childBlocked.isDone()) {
+ return forecastExecutionFuture;
+ } else if (executionDone) {
+ return childBlocked;
+ } else {
+ return successfulAsList(Arrays.asList(forecastExecutionFuture, childBlocked));
+ }
+ }
+
+ private boolean forecastExecutionDone() {
+ if (forecastExecutionFuture == null) {
+ return true;
+ }
+ return forecastExecutionFuture.isDone();
+ }
+
+ @Override
+ public boolean hasNext() throws Exception {
+ return !finished;
+ }
+
+ @Override
+ public TsBlock next() throws Exception {
+ if (forecastExecutionFuture == null) {
+ if (child.hasNextWithTimer()) {
+ TsBlock inputTsBlock = child.nextWithTimer();
+ if (inputTsBlock != null) {
+ appendTsBlockToBuilder(inputTsBlock);
+ }
+ } else {
+ submitForecastTask();
+ }
+ return null;
+ } else {
+ try {
+ if (!forecastExecutionFuture.isDone()) {
+ throw new IllegalStateException(
+ "The operator cannot continue until the forecast execution is done.");
+ }
+
+ TForecastResp forecastResp = forecastExecutionFuture.get();
+ if (forecastResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ String message =
+ String.format(
+ "Error occurred while executing forecast: %s",
+ forecastResp.getStatus().getMessage());
+ throw new ModelInferenceProcessException(message);
+ }
+
+ finished = true;
+ TsBlock resultTsBlock =
+ new TsBlockSerde().deserialize(forecastResp.bufferForForecastResult());
+ resultTsBlock = modifyTimeColumn(resultTsBlock);
+ return resultTsBlock;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new ModelInferenceProcessException(e.getMessage());
+ } catch (ExecutionException e) {
+ throw new ModelInferenceProcessException(e.getMessage());
+ }
+ }
+ }
+
+ private TsBlock modifyTimeColumn(TsBlock resultTsBlock) {
+ TsBlockBuilder newTsBlockBuilder = TsBlockBuilder.createWithOnlyTimeColumn();
+ TimeColumnBuilder timeColumnBuilder = newTsBlockBuilder.getTimeColumnBuilder();
+ for (int i = 0; i < resultTsBlock.getPositionCount(); i++) {
+ timeColumnBuilder.writeLong(
+ TimestampPrecisionUtils.convertToCurrPrecision(
+ resultTsBlock.getTimeByIndex(i), TimeUnit.NANOSECONDS));
+ newTsBlockBuilder.declarePosition();
+ }
+ return newTsBlockBuilder.build().appendValueColumns(resultTsBlock.getValueColumns());
+ }
+
+ private void appendTsBlockToBuilder(TsBlock inputTsBlock) {
+ TimeColumnBuilder timeColumnBuilder = inputTsBlockBuilder.getTimeColumnBuilder();
+ ColumnBuilder[] columnBuilders = inputTsBlockBuilder.getValueColumnBuilders();
+
+ for (int i = 0; i < inputTsBlock.getPositionCount(); i++) {
+ timeColumnBuilder.writeLong(inputTsBlock.getTimeByIndex(i));
+ for (int columnIndex = 0; columnIndex < inputTsBlock.getValueColumnCount(); columnIndex++) {
+ columnBuilders[columnIndex].write(inputTsBlock.getColumn(columnIndex), i);
+ }
+ inputTsBlockBuilder.declarePosition();
+ }
+ }
+
+ private void submitForecastTask() {
+
+ TsBlock inputTsBlock = inputTsBlockBuilder.build();
+ inputTsBlock.reverse();
+
+ forecastExecutionFuture =
+ Futures.submit(
+ () -> {
+ try (MLNodeClient client =
+ MLNodeClientManager.getInstance().borrowClient(MLNodeInfo.endPoint)) {
+ return client.forecast(
+ modelPath,
+ inputTsBlock,
+ inputTypeList,
+ inputColumnNameList,
+ expectedPredictLength);
+ } catch (Exception e) {
+ throw new ModelInferenceProcessException(e.getMessage());
+ }
+ },
+ modelInferenceExecutor);
+ }
+
+ @Override
+ public boolean isFinished() throws Exception {
+ return finished;
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (forecastExecutionFuture != null) {
+ forecastExecutionFuture.cancel(true);
+ }
+ child.close();
+ }
+
+ @Override
+ public long calculateMaxPeekMemory() {
+ return maxReturnSize + maxRetainedSize + child.calculateMaxPeekMemory();
+ }
+
+ @Override
+ public long calculateMaxReturnSize() {
+ return maxReturnSize;
+ }
+
+ @Override
+ public long calculateRetainedSizeAfterCallingNext() {
+ return maxRetainedSize + child.calculateRetainedSizeAfterCallingNext();
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java
index c49ee5a0ee6..7aae5c96fac 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java
@@ -38,6 +38,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByParame
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByTimeParameter;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.IntoPathDescriptor;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.OrderByParameter;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ModelInferenceDescriptor;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
import org.apache.iotdb.db.queryengine.plan.statement.component.SortItem;
@@ -215,6 +216,8 @@ public class Analysis {
// indicate whether the Nodes produce source data are VirtualSourceNodes
private boolean isVirtualSource = false;
+ private ModelInferenceDescriptor modelInferenceDescriptor;
+
/////////////////////////////////////////////////////////////////////////////////////////////////
// SELECT INTO Analysis
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -728,6 +731,14 @@ public class Analysis {
this.timeseriesOrderingForLastQuery = timeseriesOrderingForLastQuery;
}
+ public ModelInferenceDescriptor getModelInferenceDescriptor() {
+ return modelInferenceDescriptor;
+ }
+
+ public void setModelInferenceDescriptor(ModelInferenceDescriptor modelInferenceDescriptor) {
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ }
+
public Map<String, List<String>> getOutputDeviceToQueriedDevicesMap() {
return outputDeviceToQueriedDevicesMap;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java
index 63ba89b57db..c572d90a145 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java
@@ -26,6 +26,8 @@ import org.apache.iotdb.commons.conf.IoTDBConstant;
import org.apache.iotdb.commons.exception.IllegalPathException;
import org.apache.iotdb.commons.exception.IoTDBException;
import org.apache.iotdb.commons.exception.MetadataException;
+import org.apache.iotdb.commons.model.ForecastModeInformation;
+import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.commons.partition.DataPartition;
import org.apache.iotdb.commons.partition.DataPartitionQueryParam;
import org.apache.iotdb.commons.partition.SchemaNodeManagementPartition;
@@ -37,6 +39,7 @@ import org.apache.iotdb.commons.schema.SchemaConstant;
import org.apache.iotdb.commons.schema.view.LogicalViewSchema;
import org.apache.iotdb.commons.schema.view.viewExpression.ViewExpression;
import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
@@ -75,6 +78,8 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByTimePa
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByVariationParameter;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.IntoPathDescriptor;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.OrderByParameter;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ForecastModelInferenceDescriptor;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ModelInferenceDescriptor;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.StatementNode;
import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
@@ -86,6 +91,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.component.GroupBySessionCo
import org.apache.iotdb.db.queryengine.plan.statement.component.GroupByTimeComponent;
import org.apache.iotdb.db.queryengine.plan.statement.component.GroupByVariationComponent;
import org.apache.iotdb.db.queryengine.plan.statement.component.IntoComponent;
+import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByComponent;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
import org.apache.iotdb.db.queryengine.plan.statement.component.ResultColumn;
import org.apache.iotdb.db.queryengine.plan.statement.component.SortItem;
@@ -176,6 +182,7 @@ import static org.apache.iotdb.commons.conf.IoTDBConstant.ALLOWED_SCHEMA_PROPS;
import static org.apache.iotdb.commons.conf.IoTDBConstant.DEADBAND;
import static org.apache.iotdb.commons.conf.IoTDBConstant.LOSS;
import static org.apache.iotdb.commons.conf.IoTDBConstant.ONE_LEVEL_PATH_WILDCARD;
+import static org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction.FORECAST;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.DEVICE;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.ENDTIME;
import static org.apache.iotdb.db.queryengine.metric.QueryPlanCostMetricSet.PARTITION_FETCHER;
@@ -187,6 +194,7 @@ import static org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer.no
import static org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer.searchAggregationExpressions;
import static org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer.searchSourceExpressions;
import static org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer.toLowerCaseExpression;
+import static org.apache.iotdb.db.queryengine.plan.analyze.ExpressionTypeAnalyzer.analyzeExpression;
import static org.apache.iotdb.db.queryengine.plan.analyze.SelectIntoUtils.constructTargetDevice;
import static org.apache.iotdb.db.queryengine.plan.analyze.SelectIntoUtils.constructTargetMeasurement;
import static org.apache.iotdb.db.queryengine.plan.analyze.SelectIntoUtils.constructTargetPath;
@@ -194,6 +202,7 @@ import static org.apache.iotdb.db.queryengine.plan.optimization.LimitOffsetPushD
import static org.apache.iotdb.db.queryengine.plan.optimization.LimitOffsetPushDown.pushDownLimitOffsetInGroupByTimeForDevice;
import static org.apache.iotdb.db.schemaengine.schemaregion.view.visitor.GetSourcePathsVisitor.getSourcePaths;
import static org.apache.iotdb.db.utils.constant.SqlConstant.COUNT_TIME_HEADER;
+import static org.apache.iotdb.db.utils.constant.SqlConstant.MODEL_ID;
/** This visitor is used to analyze each type of Statement and returns the {@link Analysis}. */
public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext> {
@@ -213,6 +222,7 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
private final IPartitionFetcher partitionFetcher;
private final ISchemaFetcher schemaFetcher;
+ private final IModelFetcher modelFetcher;
private static final PerformanceOverviewMetrics PERFORMANCE_OVERVIEW_METRICS =
PerformanceOverviewMetrics.getInstance();
@@ -220,6 +230,7 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
public AnalyzeVisitor(IPartitionFetcher partitionFetcher, ISchemaFetcher schemaFetcher) {
this.partitionFetcher = partitionFetcher;
this.schemaFetcher = schemaFetcher;
+ this.modelFetcher = ModelFetcher.getInstance();
}
@Override
@@ -244,6 +255,10 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
// check for semantic errors
queryStatement.semanticCheck();
+ if (queryStatement.isModelInferenceQuery()) {
+ analyzeModelInference(analysis, queryStatement);
+ }
+
ISchemaTree schemaTree = analyzeSchema(queryStatement, analysis, context);
// If there is no leaf node in the schema tree, the query should be completed immediately
if (schemaTree.isEmpty()) {
@@ -341,6 +356,47 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
return analysis;
}
+ private void analyzeModelInference(Analysis analysis, QueryStatement queryStatement) {
+ FunctionExpression modelInferenceExpression =
+ (FunctionExpression)
+ queryStatement.getSelectComponent().getResultColumns().get(0).getExpression();
+ String modelId = modelInferenceExpression.getFunctionAttributes().get(MODEL_ID);
+
+ ModelInformation modelInformation = modelFetcher.getModelInformation(modelId);
+ if (modelInformation == null || !modelInformation.available()) {
+ throw new SemanticException("Model " + modelId + " is not available");
+ }
+
+ ModelInferenceFunction functionType =
+ ModelInferenceFunction.valueOf(modelInferenceExpression.getFunctionName().toUpperCase());
+ if (functionType == FORECAST) {
+ ForecastModelInferenceDescriptor modelInferenceDescriptor =
+ new ForecastModelInferenceDescriptor(
+ functionType, (ForecastModeInformation) modelInformation);
+ Map<String, String> modelInferenceAttributes =
+ modelInferenceExpression.getFunctionAttributes();
+ if (modelInferenceAttributes.containsKey("predict_length")) {
+ modelInferenceDescriptor.setExpectedPredictLength(
+ Integer.parseInt(modelInferenceAttributes.get("predict_length")));
+ }
+ analysis.setModelInferenceDescriptor(modelInferenceDescriptor);
+
+ List<ResultColumn> newResultColumns = new ArrayList<>();
+ for (Expression inputExpression : modelInferenceExpression.getExpressions()) {
+ newResultColumns.add(
+ new ResultColumn(inputExpression, ResultColumn.ColumnType.MODEL_INFERENCE));
+ }
+ queryStatement.getSelectComponent().setResultColumns(newResultColumns);
+
+ OrderByComponent descTimeOrder = new OrderByComponent();
+ descTimeOrder.addSortItem(new SortItem("TIME", Ordering.DESC));
+ queryStatement.setOrderByComponent(descTimeOrder);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported model inference function type " + functionType);
+ }
+ }
+
private ISchemaTree analyzeSchema(
QueryStatement queryStatement, Analysis analysis, MPPQueryContext context) {
// concat path and construct path pattern tree
@@ -1422,6 +1478,58 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
return;
}
+ if (queryStatement.isModelInferenceQuery()) {
+ List<ColumnHeader> columnHeaders = new ArrayList<>();
+
+ ModelInferenceDescriptor modelInferenceDescriptor = analysis.getModelInferenceDescriptor();
+ if (Objects.requireNonNull(modelInferenceDescriptor.getFunctionType()) == FORECAST) {
+ ForecastModelInferenceDescriptor forecastModelInferenceDescriptor =
+ (ForecastModelInferenceDescriptor) modelInferenceDescriptor;
+
+ List<TSDataType> inputTypeList = forecastModelInferenceDescriptor.getInputTypeList();
+ if (outputExpressions.size() != inputTypeList.size()) {
+ throw new SemanticException(
+ String.format(
+ "The number of input expressions does not match the number of input types [%d] when training",
+ inputTypeList.size()));
+ }
+ for (int i = 0; i < inputTypeList.size(); i++) {
+ Expression inputExpression = outputExpressions.get(i).left;
+ TSDataType inputDataType = analysis.getType(inputExpression);
+ if (inputDataType != inputTypeList.get(i)) {
+ throw new SemanticException(
+ String.format(
+ "The type of input expression [%s] does not match the type of input type [%s] when training",
+ inputDataType, inputTypeList.get(i)));
+ }
+ }
+
+ List<FunctionExpression> modelInferenceOutputExpressions = new ArrayList<>();
+ for (int predictIndex : forecastModelInferenceDescriptor.getPredictIndexList()) {
+ Expression inputExpression = outputExpressions.get(predictIndex).left;
+ FunctionExpression modelInferenceOutputExpression =
+ new FunctionExpression(
+ FORECAST.getFunctionName(),
+ forecastModelInferenceDescriptor.getOutputAttributes(),
+ Collections.singletonList(inputExpression));
+ analyzeExpression(analysis, modelInferenceOutputExpression);
+ modelInferenceOutputExpressions.add(modelInferenceOutputExpression);
+ columnHeaders.add(
+ new ColumnHeader(
+ modelInferenceOutputExpression.toString(),
+ analysis.getType(modelInferenceOutputExpression)));
+ }
+ forecastModelInferenceDescriptor.setModelInferenceOutputExpressions(
+ modelInferenceOutputExpressions);
+ } else {
+ throw new SemanticException(
+ "Unsupported model inference function type "
+ + modelInferenceDescriptor.getFunctionType());
+ }
+ analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, false));
+ return;
+ }
+
boolean isIgnoreTimestamp = queryStatement.isAggregationQuery() && !queryStatement.isGroupBy();
List<ColumnHeader> columnHeaders = new ArrayList<>();
if (queryStatement.isAlignByDevice()) {
@@ -1442,7 +1550,9 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
// For last query
private void analyzeLastOrderBy(Analysis analysis, QueryStatement queryStatement) {
- if (!queryStatement.hasOrderBy()) return;
+ if (!queryStatement.hasOrderBy()) {
+ return;
+ }
if (queryStatement.onlyOrderByTimeseries()) {
analysis.setTimeseriesOrderingForLastQuery(
@@ -1461,7 +1571,9 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
private void analyzeOrderBy(
Analysis analysis, QueryStatement queryStatement, ISchemaTree schemaTree) {
- if (!queryStatement.hasOrderByExpression()) return;
+ if (!queryStatement.hasOrderByExpression()) {
+ return;
+ }
Set<Expression> orderByExpressions = new LinkedHashSet<>();
for (Expression expressionForItem : queryStatement.getExpressionSortItemList()) {
@@ -1491,7 +1603,7 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext>
}
private TSDataType analyzeExpressionType(Analysis analysis, Expression expression) {
- return ExpressionTypeAnalyzer.analyzeExpression(analysis, expression);
+ return analyzeExpression(analysis, expression);
}
private void analyzeDeviceToGroupBy(
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionAnalyzer.java
index 28b86e4a3d4..841f9254ab2 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionAnalyzer.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionAnalyzer.java
@@ -49,6 +49,7 @@ import org.apache.iotdb.db.queryengine.plan.expression.visitor.CollectSourceExpr
import org.apache.iotdb.db.queryengine.plan.expression.visitor.ExpressionNormalizeVisitor;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.GetMeasurementExpressionVisitor;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.LowercaseNormalizeVisitor;
+import org.apache.iotdb.db.queryengine.plan.expression.visitor.RemoveRootPrefixVisitor;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.ReplaceRawPathWithGroupedPathVisitor;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.cartesian.BindSchemaForExpressionVisitor;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.cartesian.BindSchemaForPredicateVisitor;
@@ -196,6 +197,8 @@ public class ExpressionAnalyzer {
}
}
return ResultColumn.ColumnType.AGGREGATION;
+ } else if (((FunctionExpression) expression).isModelInferenceFunction()) {
+ return ResultColumn.ColumnType.MODEL_INFERENCE;
} else {
ResultColumn.ColumnType checkedType = null;
int lastCheckedIndex = 0;
@@ -846,4 +849,8 @@ public class ExpressionAnalyzer {
throw new UnknownExpressionTypeException(expression.getExpressionType());
}
}
+
+ public static Expression removeRootPrefix(Expression expression) {
+ return new RemoveRootPrefixVisitor().process(expression, null);
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionTypeAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionTypeAnalyzer.java
index 41cd1fbe882..18eb59894b7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionTypeAnalyzer.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionTypeAnalyzer.java
@@ -269,6 +269,8 @@ public class ExpressionTypeAnalyzer {
functionExpression,
TypeInferenceUtils.getBuiltInScalarFunctionDataType(
functionExpression, expressionTypes.get(NodeRef.of(inputExpressions.get(0)))));
+ } else if (functionExpression.isModelInferenceFunction()) {
+ return setExpressionType(functionExpression, TSDataType.DOUBLE);
} else {
return setExpressionType(
functionExpression,
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java
similarity index 75%
copy from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
copy to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java
index 304d209817f..1f0221b5940 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java
@@ -17,11 +17,11 @@
* under the License.
*/
-package org.apache.iotdb.confignode.procedure.state.model;
+package org.apache.iotdb.db.queryengine.plan.analyze;
-public enum CreateModelState {
- INIT,
- VALIDATED,
- CONFIG_NODE_ACTIVE,
- ML_NODE_ACTIVE
+import org.apache.iotdb.commons.model.ModelInformation;
+
+public interface IModelFetcher {
+ /** Get model information by model id from configNode. */
+ ModelInformation getModelInformation(String modelId);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ModelFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ModelFetcher.java
new file mode 100644
index 00000000000..d829c51c0d0
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ModelFetcher.java
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.analyze;
+
+import org.apache.iotdb.commons.client.IClientManager;
+import org.apache.iotdb.commons.client.exception.ClientManagerException;
+import org.apache.iotdb.commons.consensus.ConfigRegionId;
+import org.apache.iotdb.commons.model.ModelInformation;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
+import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
+import org.apache.iotdb.db.exception.sql.StatementAnalyzeException;
+import org.apache.iotdb.db.protocol.client.ConfigNodeClient;
+import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
+import org.apache.iotdb.db.protocol.client.ConfigNodeInfo;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.apache.thrift.TException;
+
+public class ModelFetcher implements IModelFetcher {
+
+ private final IClientManager<ConfigRegionId, ConfigNodeClient> configNodeClientManager =
+ ConfigNodeClientManager.getInstance();
+
+ private static final class ModelFetcherHolder {
+
+ private static final ModelFetcher INSTANCE = new ModelFetcher();
+
+ private ModelFetcherHolder() {}
+ }
+
+ public static ModelFetcher getInstance() {
+ return ModelFetcher.ModelFetcherHolder.INSTANCE;
+ }
+
+ private ModelFetcher() {}
+
+ @Override
+ public ModelInformation getModelInformation(String modelId) {
+ try (ConfigNodeClient client =
+ configNodeClientManager.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TGetModelInfoResp getModelInfoResp = client.getModelInfo(new TGetModelInfoReq(modelId));
+ if (getModelInfoResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ if (getModelInfoResp.modelInfo != null) {
+ return ModelInformation.deserialize(getModelInfoResp.modelInfo);
+ } else {
+ return null;
+ }
+ } else {
+ throw new StatementAnalyzeException(
+ "An error occurred when executing getModelInformation():"
+ + getModelInfoResp.getStatus().getMessage());
+ }
+ } catch (ClientManagerException | TException e) {
+ throw new StatementAnalyzeException(
+ "An error occurred when executing getModelInformation():" + e.getMessage());
+ }
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigTaskVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigTaskVisitor.java
index 57b6a3cbad9..39382367fa5 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigTaskVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigTaskVisitor.java
@@ -53,7 +53,7 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.UnSetTTLTa
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.CreateModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.DropModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowModelsTask;
-import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowTrailsTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowTrialsTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.AlterSchemaTemplateTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.CreateSchemaTemplateTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.DeactivateSchemaTemplateTask;
@@ -118,7 +118,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.UnSetTTLStatement
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.DropModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
-import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrailsStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrialsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.DropPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.ShowPipesStatement;
@@ -522,8 +522,8 @@ public class ConfigTaskVisitor
}
@Override
- public IConfigTask visitShowTrails(ShowTrailsStatement showTrailsStatement, TaskContext context) {
- return new ShowTrailsTask(showTrailsStatement.getModelId());
+ public IConfigTask visitShowTrials(ShowTrialsStatement showTrialsStatement, TaskContext context) {
+ return new ShowTrialsTask(showTrialsStatement.getModelId());
}
public static class TaskContext {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
index 1f25a869d04..934d41568e4 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
@@ -90,8 +90,8 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowTrailResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowTrialResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp;
import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp;
import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp;
@@ -110,6 +110,7 @@ import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.analyze.Analysis;
import org.apache.iotdb.db.queryengine.plan.analyze.Analyzer;
import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
+import org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer;
import org.apache.iotdb.db.queryengine.plan.analyze.schema.ClusterSchemaFetcher;
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
@@ -131,7 +132,7 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowTTLTas
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowTriggersTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowVariablesTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowModelsTask;
-import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowTrailsTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowTrialsTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.ShowNodesInSchemaTemplateTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.ShowPathSetTemplateTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.template.ShowSchemaTemplateTask;
@@ -140,6 +141,11 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowSpace
import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowThrottleQuotaTask;
import org.apache.iotdb.db.queryengine.plan.expression.Expression;
import org.apache.iotdb.db.queryengine.plan.expression.visitor.TransformToViewExpressionVisitor;
+import org.apache.iotdb.db.queryengine.plan.statement.component.FromComponent;
+import org.apache.iotdb.db.queryengine.plan.statement.component.ResultColumn;
+import org.apache.iotdb.db.queryengine.plan.statement.component.SelectComponent;
+import org.apache.iotdb.db.queryengine.plan.statement.component.WhereCondition;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.CountDatabaseStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.CountTimeSlotListStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.CreateContinuousQueryStatement;
@@ -196,6 +202,8 @@ import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.trigger.api.Trigger;
import org.apache.iotdb.trigger.api.enums.FailureStrategy;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
+import org.apache.iotdb.tsfile.utils.Pair;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
import org.apache.iotdb.udf.api.UDTF;
@@ -226,7 +234,14 @@ import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.INPUT_TYPE_LIST;
+import static org.apache.iotdb.commons.model.ForecastModeInformation.PREDICT_INDEX_LIST;
+import static org.apache.iotdb.commons.schema.SchemaConstant.ROOT;
import static org.apache.iotdb.db.protocol.client.ConfigNodeClient.MSG_RECONNECTION_FAIL;
+import static org.apache.iotdb.db.queryengine.plan.expression.ExpressionFactory.and;
+import static org.apache.iotdb.db.queryengine.plan.expression.ExpressionFactory.longValue;
+import static org.apache.iotdb.db.queryengine.plan.expression.ExpressionFactory.lte;
+import static org.apache.iotdb.db.queryengine.plan.expression.ExpressionFactory.time;
public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
@@ -2085,57 +2100,74 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
}
@Override
- public SettableFuture<ConfigTaskResult> setSpaceQuota(
- SetSpaceQuotaStatement setSpaceQuotaStatement) {
+ public SettableFuture<ConfigTaskResult> createModel(CreateModelStatement createModelStatement) {
+ createModelStatement.semanticCheck();
+
+ QueryStatement datasetStatement = createModelStatement.getDatasetStatement();
+ Analyzer analyzer = Analyzer.getAnalyzer();
+ Analysis analysis = analyzer.analyze(datasetStatement);
+
SettableFuture<ConfigTaskResult> future = SettableFuture.create();
- TSStatus tsStatus = new TSStatus();
- TSetSpaceQuotaReq req = new TSetSpaceQuotaReq();
- req.setDatabase(setSpaceQuotaStatement.getPrefixPathList());
- TSpaceQuota spaceQuota = new TSpaceQuota();
- spaceQuota.setDeviceNum(setSpaceQuotaStatement.getDeviceNum());
- spaceQuota.setTimeserieNum(setSpaceQuotaStatement.getTimeSeriesNum());
- spaceQuota.setDiskSize(setSpaceQuotaStatement.getDiskSize());
- req.setSpaceLimit(spaceQuota);
- try (ConfigNodeClient client =
- CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- // Send request to some API server
- tsStatus = client.setSpaceQuota(req);
- } catch (Exception e) {
- future.setException(e);
+ if (analysis.getOutputExpressions() == null) {
+ future.setException(
+ new IoTDBException(
+ "The query result of the sql in CREATE MODEL are null.",
+ TSStatusCode.ILLEGAL_PARAMETER.getStatusCode()));
+ return future;
}
- if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
- } else {
- future.setException(new IoTDBException(tsStatus.message, tsStatus.code));
+
+ SelectComponent formattedSelect =
+ new SelectComponent(datasetStatement.getSelectComponent().getZoneId());
+ List<TSDataType> inputTypeList = new ArrayList<>();
+ for (Expression outputExpression :
+ analysis.getOutputExpressions().stream().map(Pair::getLeft).collect(Collectors.toList())) {
+ ResultColumn.ColumnType columnType =
+ ExpressionAnalyzer.identifyOutputColumnType(outputExpression, true);
+ formattedSelect.addResultColumn(
+ new ResultColumn(ExpressionAnalyzer.removeRootPrefix(outputExpression), columnType));
+ inputTypeList.add(analysis.getType(outputExpression));
}
- return future;
- }
+ datasetStatement.setSelectComponent(formattedSelect);
- @Override
- public SettableFuture<ConfigTaskResult> createModel(CreateModelStatement createModelStatement) {
- createModelStatement.semanticCheck();
+ FromComponent fromRoot = new FromComponent();
+ fromRoot.addPrefixPath(new PartialPath(ROOT, false));
+ datasetStatement.setFromComponent(fromRoot);
- Analyzer analyzer = Analyzer.getAnalyzer();
- Analysis analysis = analyzer.analyze(createModelStatement.getQueryStatement());
+ WhereCondition endNowWhereCondition = new WhereCondition();
+ Expression endNowExpression = lte(time(), longValue(System.currentTimeMillis()));
+ if (!datasetStatement.hasWhere()) {
+ endNowWhereCondition.setPredicate(endNowExpression);
+ } else {
+ endNowWhereCondition.setPredicate(
+ and(
+ ExpressionAnalyzer.removeRootPrefix(
+ datasetStatement.getWhereCondition().getPredicate()),
+ endNowExpression));
+ }
+ datasetStatement.setWhereCondition(endNowWhereCondition);
+
+ String datesetFetchSQL = datasetStatement.constructFormattedSQL();
- List<String> queryExpressions = new ArrayList<>();
- for (Expression expression : analysis.getSelectExpressions()) {
- queryExpressions.add(expression.getExpressionString());
+ Map<String, String> options = createModelStatement.getOptions();
+ options.put(INPUT_TYPE_LIST, Arrays.toString(inputTypeList.toArray()));
+ if (!options.containsKey(PREDICT_INDEX_LIST)) {
+ List<Integer> predictIndexList = new ArrayList<>();
+ for (int i = 0; i < inputTypeList.size(); i++) {
+ predictIndexList.add(i);
+ }
+ options.put(PREDICT_INDEX_LIST, Arrays.toString(predictIndexList.toArray()));
}
- Expression whereExpression = analysis.getWhereExpression();
- String queryFilter = whereExpression == null ? null : whereExpression.getExpressionString();
- SettableFuture<ConfigTaskResult> future = SettableFuture.create();
try (ConfigNodeClient client =
CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- TCreateModelReq createModelReq = new TCreateModelReq();
- createModelReq.setModelId(createModelStatement.getModelId());
- createModelReq.setModelTask(createModelStatement.getModelTask());
- createModelReq.setModelType(createModelStatement.getModelType());
- createModelReq.setIsAuto(createModelStatement.isAuto());
- createModelReq.setQueryExpressions(queryExpressions);
- createModelReq.setQueryFilter(queryFilter);
- createModelReq.setModelConfigs(createModelStatement.getAttributes());
+ TCreateModelReq createModelReq =
+ new TCreateModelReq(
+ createModelStatement.getModelId(),
+ createModelStatement.getTaskType(),
+ createModelStatement.getOptions(),
+ createModelStatement.getHyperparameters(),
+ datesetFetchSQL);
+
final TSStatus executionStatus = client.createModel(createModelReq);
if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != executionStatus.getCode()) {
LOGGER.warn(
@@ -2153,6 +2185,89 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
return future;
}
+ @Override
+ public SettableFuture<ConfigTaskResult> dropModel(String modelId) {
+ SettableFuture<ConfigTaskResult> future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ final TSStatus executionStatus = client.dropModel(new TDropModelReq(modelId));
+ if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != executionStatus.getCode()) {
+ LOGGER.warn("[{}] Failed to drop model {}.", executionStatus, modelId);
+ future.setException(new IoTDBException(executionStatus.message, executionStatus.code));
+ } else {
+ future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
+ }
+ } catch (ClientManagerException | TException e) {
+ future.setException(e);
+ }
+ return future;
+ }
+
+ @Override
+ public SettableFuture<ConfigTaskResult> showModels() {
+ SettableFuture<ConfigTaskResult> future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TShowModelResp showModelResp = client.showModel(new TShowModelReq());
+ if (showModelResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ future.setException(
+ new IoTDBException(showModelResp.getStatus().message, showModelResp.getStatus().code));
+ return future;
+ }
+ // convert model info list and buildTsBlock
+ ShowModelsTask.buildTsBlock(showModelResp.getModelInfoList(), future);
+ } catch (ClientManagerException | TException e) {
+ future.setException(e);
+ }
+ return future;
+ }
+
+ @Override
+ public SettableFuture<ConfigTaskResult> showTrials(String modelId) {
+ SettableFuture<ConfigTaskResult> future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TShowTrialResp showTrialResp = client.showTrial(new TShowTrialReq(modelId));
+ if (showTrialResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ future.setException(
+ new IoTDBException(showTrialResp.getStatus().message, showTrialResp.getStatus().code));
+ return future;
+ }
+ // convert trail info list and buildTsBlock
+ ShowTrialsTask.buildTsBlock(showTrialResp.getTrialInfoList(), future);
+ } catch (ClientManagerException | TException e) {
+ future.setException(e);
+ }
+ return future;
+ }
+
+ @Override
+ public SettableFuture<ConfigTaskResult> setSpaceQuota(
+ SetSpaceQuotaStatement setSpaceQuotaStatement) {
+ SettableFuture<ConfigTaskResult> future = SettableFuture.create();
+ TSStatus tsStatus = new TSStatus();
+ TSetSpaceQuotaReq req = new TSetSpaceQuotaReq();
+ req.setDatabase(setSpaceQuotaStatement.getPrefixPathList());
+ TSpaceQuota spaceQuota = new TSpaceQuota();
+ spaceQuota.setDeviceNum(setSpaceQuotaStatement.getDeviceNum());
+ spaceQuota.setTimeserieNum(setSpaceQuotaStatement.getTimeSeriesNum());
+ spaceQuota.setDiskSize(setSpaceQuotaStatement.getDiskSize());
+ req.setSpaceLimit(spaceQuota);
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ // Send request to some API server
+ tsStatus = client.setSpaceQuota(req);
+ } catch (Exception e) {
+ future.setException(e);
+ }
+ if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
+ } else {
+ future.setException(new IoTDBException(tsStatus.message, tsStatus.code));
+ }
+ return future;
+ }
+
@Override
public SettableFuture<ConfigTaskResult> showSpaceQuota(
ShowSpaceQuotaStatement showSpaceQuotaStatement) {
@@ -2235,43 +2350,6 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
return throttleQuotaResp;
}
- @Override
- public SettableFuture<ConfigTaskResult> dropModel(String modelId) {
- SettableFuture<ConfigTaskResult> future = SettableFuture.create();
- try (ConfigNodeClient client =
- CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- final TSStatus executionStatus = client.dropModel(new TDropModelReq(modelId));
- if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != executionStatus.getCode()) {
- LOGGER.warn("[{}] Failed to drop model {}.", executionStatus, modelId);
- future.setException(new IoTDBException(executionStatus.message, executionStatus.code));
- } else {
- future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
- }
- } catch (ClientManagerException | TException e) {
- future.setException(e);
- }
- return future;
- }
-
- @Override
- public SettableFuture<ConfigTaskResult> showModels() {
- SettableFuture<ConfigTaskResult> future = SettableFuture.create();
- try (ConfigNodeClient client =
- CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- TShowModelResp showModelResp = client.showModel(new TShowModelReq());
- if (showModelResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- future.setException(
- new IoTDBException(showModelResp.getStatus().message, showModelResp.getStatus().code));
- return future;
- }
- // convert model info list and buildTsBlock
- ShowModelsTask.buildTsBlock(showModelResp.getModelInfoList(), future);
- } catch (ClientManagerException | TException e) {
- future.setException(e);
- }
- return future;
- }
-
@Override
public TSpaceQuotaResp getSpaceQuota() {
TSpaceQuotaResp spaceQuotaResp = new TSpaceQuotaResp();
@@ -2284,23 +2362,4 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
}
return spaceQuotaResp;
}
-
- @Override
- public SettableFuture<ConfigTaskResult> showTrails(String modelId) {
- SettableFuture<ConfigTaskResult> future = SettableFuture.create();
- try (ConfigNodeClient client =
- CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- TShowTrailResp showTrailResp = client.showTrail(new TShowTrailReq(modelId));
- if (showTrailResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- future.setException(
- new IoTDBException(showTrailResp.getStatus().message, showTrailResp.getStatus().code));
- return future;
- }
- // convert trail info list and buildTsBlock
- ShowTrailsTask.buildTsBlock(showTrailResp.getTrailInfoList(), future);
- } catch (ClientManagerException | TException e) {
- future.setException(e);
- }
- return future;
- }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
index 895f81cd9be..7d84aec7b4a 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
@@ -214,5 +214,5 @@ public interface IConfigTaskExecutor {
SettableFuture<ConfigTaskResult> showModels();
- SettableFuture<ConfigTaskResult> showTrails(String modelId);
+ SettableFuture<ConfigTaskResult> showTrials(String modelId);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowModelsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowModelsTask.java
index 9f34d0b052d..04e50abbd9f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowModelsTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowModelsTask.java
@@ -42,6 +42,8 @@ import java.util.stream.Collectors;
public class ShowModelsTask implements IConfigTask {
+ private static final int MODEL_INFO_COLUMN_NUM = 5;
+
public ShowModelsTask() {
// do nothing
}
@@ -65,7 +67,6 @@ public class ShowModelsTask implements IConfigTask {
String modelType = ReadWriteIOUtils.readString(modelInfo);
String queryBody = ReadWriteIOUtils.readString(modelInfo);
String trainingState = ReadWriteIOUtils.readString(modelInfo);
- String modelPath = ReadWriteIOUtils.readString(modelInfo);
int listSize = ReadWriteIOUtils.readInt(modelInfo);
List<String> modelHyperparameter = new ArrayList<>();
@@ -79,16 +80,17 @@ public class ShowModelsTask implements IConfigTask {
builder.getColumnBuilder(2).writeBinary(Binary.valueOf(modelType));
builder.getColumnBuilder(3).writeBinary(Binary.valueOf(queryBody));
builder.getColumnBuilder(4).writeBinary(Binary.valueOf(trainingState));
- builder.getColumnBuilder(5).writeBinary(Binary.valueOf(modelPath));
- builder.getColumnBuilder(6).writeBinary(Binary.valueOf(modelHyperparameter.get(0)));
+ builder.getColumnBuilder(5).writeBinary(Binary.valueOf(modelHyperparameter.get(0)));
builder.declarePosition();
for (int i = 1; i < listSize; i++) {
builder.getTimeColumnBuilder().writeLong(0L);
- for (int columnIndex = 0; columnIndex <= 5; columnIndex++) {
+ for (int columnIndex = 0; columnIndex <= MODEL_INFO_COLUMN_NUM - 1; columnIndex++) {
builder.getColumnBuilder(columnIndex).writeBinary(Binary.valueOf(""));
}
- builder.getColumnBuilder(6).writeBinary(Binary.valueOf(modelHyperparameter.get(i)));
+ builder
+ .getColumnBuilder(MODEL_INFO_COLUMN_NUM)
+ .writeBinary(Binary.valueOf(modelHyperparameter.get(i)));
builder.declarePosition();
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrailsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrialsTask.java
similarity index 94%
rename from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrailsTask.java
rename to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrialsTask.java
index 07cc2503292..0bc5160acb0 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrailsTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowTrialsTask.java
@@ -40,24 +40,24 @@ import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
-public class ShowTrailsTask implements IConfigTask {
+public class ShowTrialsTask implements IConfigTask {
private final String modelId;
- public ShowTrailsTask(String modelId) {
+ public ShowTrialsTask(String modelId) {
this.modelId = modelId;
}
@Override
public ListenableFuture<ConfigTaskResult> execute(IConfigTaskExecutor configTaskExecutor)
throws InterruptedException {
- return configTaskExecutor.showTrails(modelId);
+ return configTaskExecutor.showTrials(modelId);
}
public static void buildTsBlock(
List<ByteBuffer> trailInfoList, SettableFuture<ConfigTaskResult> future) {
List<TSDataType> outputDataTypes =
- ColumnHeaderConstant.showTrailsColumnHeaders.stream()
+ ColumnHeaderConstant.showTrialsColumnHeaders.stream()
.map(ColumnHeader::getColumnType)
.collect(Collectors.toList());
TsBlockBuilder builder = new TsBlockBuilder(outputDataTypes);
@@ -84,7 +84,7 @@ public class ShowTrailsTask implements IConfigTask {
builder.declarePosition();
}
}
- DatasetHeader datasetHeader = DatasetHeaderFactory.getShowTrailsHeader();
+ DatasetHeader datasetHeader = DatasetHeaderFactory.getShowTrialsHeader();
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS, builder.build(), datasetHeader));
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/ExpressionFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/ExpressionFactory.java
index 741e7c4ccc0..cb34ec72682 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/ExpressionFactory.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/ExpressionFactory.java
@@ -25,6 +25,7 @@ import org.apache.iotdb.commons.path.PartialPath;
import org.apache.iotdb.db.queryengine.plan.expression.binary.AdditionExpression;
import org.apache.iotdb.db.queryengine.plan.expression.binary.EqualToExpression;
import org.apache.iotdb.db.queryengine.plan.expression.binary.GreaterThanExpression;
+import org.apache.iotdb.db.queryengine.plan.expression.binary.LessEqualExpression;
import org.apache.iotdb.db.queryengine.plan.expression.binary.LessThanExpression;
import org.apache.iotdb.db.queryengine.plan.expression.binary.LogicAndExpression;
import org.apache.iotdb.db.queryengine.plan.expression.binary.LogicOrExpression;
@@ -66,6 +67,10 @@ public class ExpressionFactory {
return new ConstantOperand(TSDataType.INT32, valueString);
}
+ public static ConstantOperand longValue(long value) {
+ return new ConstantOperand(TSDataType.INT64, String.valueOf(value));
+ }
+
public static TimestampOperand time() {
return new TimestampOperand();
}
@@ -123,6 +128,10 @@ public class ExpressionFactory {
return new LessThanExpression(leftExpression, rightExpression);
}
+ public static LessEqualExpression lte(Expression leftExpression, Expression rightExpression) {
+ return new LessEqualExpression(leftExpression, rightExpression);
+ }
+
public static EqualToExpression eq(Expression leftExpression, Expression rightExpression) {
return new EqualToExpression(leftExpression, rightExpression);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionExpression.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionExpression.java
index 17b628f1cb5..29615c489ae 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionExpression.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionExpression.java
@@ -23,6 +23,7 @@ import org.apache.iotdb.commons.conf.IoTDBConstant;
import org.apache.iotdb.commons.path.PartialPath;
import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
import org.apache.iotdb.commons.udf.builtin.BuiltinScalarFunction;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
import org.apache.iotdb.db.queryengine.common.NodeRef;
import org.apache.iotdb.db.queryengine.plan.expression.Expression;
import org.apache.iotdb.db.queryengine.plan.expression.ExpressionType;
@@ -119,6 +120,8 @@ public class FunctionExpression extends Expression {
functionType = FunctionType.AGGREGATION_FUNCTION;
} else if (BuiltinScalarFunction.getNativeFunctionNames().contains(lowerCaseFunctionName)) {
functionType = FunctionType.BUILT_IN_SCALAR_FUNCTION;
+ } else if (ModelInferenceFunction.getNativeFunctionNames().contains(functionName)) {
+ functionType = FunctionType.MODEL_INFERENCE_FUNCTION;
} else {
functionType = FunctionType.UDF;
}
@@ -139,6 +142,13 @@ public class FunctionExpression extends Expression {
return functionType == FunctionType.BUILT_IN_SCALAR_FUNCTION;
}
+ public boolean isModelInferenceFunction() {
+ if (functionType == null) {
+ initializeFunctionType();
+ }
+ return functionType == FunctionType.MODEL_INFERENCE_FUNCTION;
+ }
+
@Override
public boolean isConstantOperandInternal() {
if (isConstantOperandCache == null) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
index 336bbcce0bb..e1d64f29680 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
@@ -23,5 +23,6 @@ package org.apache.iotdb.db.queryengine.plan.expression.multi;
public enum FunctionType {
AGGREGATION_FUNCTION,
BUILT_IN_SCALAR_FUNCTION,
- UDF
+ UDF,
+ MODEL_INFERENCE_FUNCTION
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/visitor/RemoveRootPrefixVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/visitor/RemoveRootPrefixVisitor.java
new file mode 100644
index 00000000000..ada5d3c89ca
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/visitor/RemoveRootPrefixVisitor.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.expression.visitor;
+
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.db.queryengine.plan.expression.Expression;
+import org.apache.iotdb.db.queryengine.plan.expression.leaf.TimeSeriesOperand;
+
+import static org.apache.iotdb.commons.schema.SchemaConstant.ROOT;
+
+public class RemoveRootPrefixVisitor extends ReconstructVisitor<Void> {
+
+ @Override
+ public Expression visitTimeSeriesOperand(TimeSeriesOperand timeSeriesOperand, Void context) {
+ PartialPath rawPath = timeSeriesOperand.getPath();
+ if (!rawPath.startsWith(ROOT)) {
+ return timeSeriesOperand;
+ }
+
+ String[] rawPartialNodes = rawPath.getNodes();
+
+ int newPartialNodesLength = rawPartialNodes.length - 1;
+ String[] newPartialNodes = new String[newPartialNodesLength];
+ System.arraycopy(rawPartialNodes, 1, newPartialNodes, 0, newPartialNodesLength);
+
+ return new TimeSeriesOperand(new PartialPath(newPartialNodes));
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
index 58cd8d43bb4..a9e0b9e0efc 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
@@ -159,7 +159,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.UnSetTTLStatement
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.DropModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
-import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrailsStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrialsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.DropPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.ShowPipesStatement;
@@ -1223,20 +1223,46 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
public Statement visitCreateModel(IoTDBSqlParser.CreateModelContext ctx) {
CreateModelStatement createModelStatement = new CreateModelStatement();
createModelStatement.setModelId(parseIdentifier(ctx.modelId.getText()));
- createModelStatement.setAuto(ctx.AUTO() != null);
- Map<String, String> attributes = new HashMap<>();
+ Map<String, String> modelOptions = new HashMap<>();
for (IoTDBSqlParser.AttributePairContext attribute : ctx.attributePair()) {
- attributes.put(
+ modelOptions.put(
parseAttributeKey(attribute.key).toLowerCase(), parseAttributeValue(attribute.value));
}
- createModelStatement.setAttributes(attributes);
+ createModelStatement.setOptions(modelOptions);
+
+ Map<String, String> hyperParameters = new HashMap<>();
+ for (IoTDBSqlParser.HparamPairContext attribute : ctx.hparamPair()) {
+ hyperParameters.put(
+ parseAttributeKey(attribute.hparamKey).toLowerCase(),
+ parseHparamValue(attribute.hparamValue()));
+ }
+ createModelStatement.setHyperparameters(hyperParameters);
- createModelStatement.setQueryStatement(
+ createModelStatement.setDatasetStatement(
(QueryStatement) visitSelectStatement(ctx.selectStatement()));
return createModelStatement;
}
+ private String parseHparamValue(IoTDBSqlParser.HparamValueContext ctx) {
+ if (ctx.attributeValue() != null) {
+ return parseAttributeValue(ctx.attributeValue());
+ } else if (ctx.hparamRange() != null) {
+ return parseAttributeValue(ctx.hparamRange().hparamRangeStart)
+ + ","
+ + parseAttributeValue(ctx.hparamRange().hparamRangeEnd);
+ } else if (ctx.hparamCandidates() != null) {
+ List<String> candidates = new ArrayList<>();
+ for (IoTDBSqlParser.AttributeValueContext candidate :
+ ctx.hparamCandidates().attributeValue()) {
+ candidates.add(parseAttributeValue(candidate));
+ }
+ return Arrays.toString(candidates.toArray());
+ } else {
+ throw new IllegalArgumentException();
+ }
+ }
+
// Drop Model =====================================================================
@Override
public Statement visitDropModel(IoTDBSqlParser.DropModelContext ctx) {
@@ -1251,8 +1277,8 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
// Show Trails =====================================================================
@Override
- public Statement visitShowTrails(IoTDBSqlParser.ShowTrailsContext ctx) {
- return new ShowTrailsStatement(parseIdentifier(ctx.modelId.getText()));
+ public Statement visitShowTrials(IoTDBSqlParser.ShowTrialsContext ctx) {
+ return new ShowTrialsStatement(parseIdentifier(ctx.modelId.getText()));
}
/** Data Manipulation Language (DML). */
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java
index 74cb4fb8aa2..fdb48f84d43 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java
@@ -111,7 +111,6 @@ import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
-import java.util.Objects;
import static org.apache.iotdb.commons.conf.IoTDBConstant.MULTI_LEVEL_PATH_WILDCARD;
import static org.apache.iotdb.db.protocol.thrift.impl.MLNodeRPCServiceImpl.ML_METRICS_PATH_PREFIX;
@@ -865,53 +864,8 @@ public class StatementGenerator {
return insertRowStatement;
}
- public static Statement createStatement(TFetchTimeseriesReq fetchTimeseriesReq, ZoneId zoneId)
- throws IllegalPathException {
- QueryStatement queryStatement = new QueryStatement();
-
- FromComponent fromComponent = new FromComponent();
- for (String pathStr : fetchTimeseriesReq.getQueryExpressions()) {
- PartialPath path = new PartialPath(pathStr);
- fromComponent.addPrefixPath(path);
- }
- queryStatement.setFromComponent(fromComponent);
-
- SelectComponent selectComponent = new SelectComponent(zoneId);
- selectComponent.addResultColumn(
- new ResultColumn(
- new TimeSeriesOperand(new PartialPath("", false)), ResultColumn.ColumnType.RAW));
- queryStatement.setSelectComponent(selectComponent);
-
- if (fetchTimeseriesReq.isSetQueryFilter()) {
- WhereCondition whereCondition = new WhereCondition();
- String queryFilter = fetchTimeseriesReq.getQueryFilter();
- String[] times = queryFilter.split(",");
- int predictNum = 0;
- LessThanExpression rightPredicate = null;
- GreaterEqualExpression leftPredicate = null;
- if (!Objects.equals(times[0], "-1")) {
- leftPredicate =
- new GreaterEqualExpression(
- new TimestampOperand(), new ConstantOperand(TSDataType.INT64, times[0]));
- predictNum += 1;
- }
- if (!Objects.equals(times[1], "-1")) {
- rightPredicate =
- new LessThanExpression(
- new TimestampOperand(), new ConstantOperand(TSDataType.INT64, times[1]));
- predictNum += 2;
- }
-
- if (predictNum == 3) {
- whereCondition.setPredicate(new LogicAndExpression(leftPredicate, rightPredicate));
- } else if (predictNum == 1) {
- whereCondition.setPredicate(leftPredicate);
- } else {
- whereCondition.setPredicate(rightPredicate);
- }
- queryStatement.setWhereCondition(whereCondition);
- }
- return queryStatement;
+ public static Statement createStatement(TFetchTimeseriesReq fetchTimeseriesReq, ZoneId zoneId) {
+ return invokeParser(fetchTimeseriesReq.getQueryBody(), zoneId);
}
public static DeleteTimeSeriesStatement createStatement(TDeleteModelMetricsReq req)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java
index 54f2ce71dbc..c8bc458a045 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java
@@ -71,6 +71,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.SortNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TimeJoinNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TransformNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml.ForecastNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedLastQueryScanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedSeriesAggregationScanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedSeriesScanNode;
@@ -88,6 +89,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByParame
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByTimeParameter;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.IntoPathDescriptor;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.OrderByParameter;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ForecastModelInferenceDescriptor;
import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByKey;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
import org.apache.iotdb.db.queryengine.plan.statement.component.SortItem;
@@ -1393,4 +1395,12 @@ public class LogicalPlanBuilder {
}
return this;
}
+
+ public LogicalPlanBuilder planForecast(
+ ForecastModelInferenceDescriptor forecastModelInferenceDescriptor) {
+ this.root =
+ new ForecastNode(
+ context.getQueryId().genPlanNodeId(), root, forecastModelInferenceDescriptor);
+ return this;
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java
index d5180aacf46..63c3e3bf8ce 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java
@@ -23,6 +23,7 @@ import org.apache.iotdb.commons.path.MeasurementPath;
import org.apache.iotdb.commons.path.PartialPath;
import org.apache.iotdb.commons.schema.view.viewExpression.ViewExpression;
import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
import org.apache.iotdb.db.queryengine.plan.analyze.Analysis;
import org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer;
@@ -52,6 +53,8 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsOf
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.PipeEnrichedInsertNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.AggregationStep;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ForecastModelInferenceDescriptor;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ModelInferenceDescriptor;
import org.apache.iotdb.db.queryengine.plan.statement.StatementNode;
import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
import org.apache.iotdb.db.queryengine.plan.statement.crud.DeleteDataStatement;
@@ -97,6 +100,7 @@ import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Set;
import static org.apache.iotdb.db.utils.constant.SqlConstant.COUNT_TIME;
@@ -214,6 +218,22 @@ public class LogicalPlanVisitor extends StatementVisitor<PlanNode, MPPQueryConte
.planOffset(queryStatement.getRowOffset())
.planLimit(queryStatement.getRowLimit());
+ if (queryStatement.isModelInferenceQuery()) {
+ ModelInferenceDescriptor modelInferenceDescriptor = analysis.getModelInferenceDescriptor();
+ if (Objects.requireNonNull(modelInferenceDescriptor.getFunctionType())
+ == ModelInferenceFunction.FORECAST) {
+ ForecastModelInferenceDescriptor forecastModelInferenceDescriptor =
+ (ForecastModelInferenceDescriptor) modelInferenceDescriptor;
+ planBuilder
+ .planLimit(forecastModelInferenceDescriptor.getModelInputLength())
+ .planForecast(forecastModelInferenceDescriptor);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported model inference function type: "
+ + modelInferenceDescriptor.getFunctionType());
+ }
+ }
+
// plan select into
if (queryStatement.isAlignByDevice()) {
planBuilder = planBuilder.planDeviceViewInto(analysis.getDeviceViewIntoPathDescriptor());
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
index 3585b301e10..2e5a4653efe 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
@@ -102,6 +102,7 @@ import org.apache.iotdb.db.queryengine.execution.operator.process.last.LastQuery
import org.apache.iotdb.db.queryengine.execution.operator.process.last.LastQueryUtil;
import org.apache.iotdb.db.queryengine.execution.operator.process.last.UpdateLastCacheOperator;
import org.apache.iotdb.db.queryengine.execution.operator.process.last.UpdateViewPathLastCacheOperator;
+import org.apache.iotdb.db.queryengine.execution.operator.process.ml.ForecastOperator;
import org.apache.iotdb.db.queryengine.execution.operator.schema.CountGroupByLevelMergeOperator;
import org.apache.iotdb.db.queryengine.execution.operator.schema.CountGroupByLevelScanOperator;
import org.apache.iotdb.db.queryengine.execution.operator.schema.CountMergeOperator;
@@ -179,6 +180,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TransformN
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryCollectNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml.ForecastNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.ShuffleSinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedLastQueryScanNode;
@@ -202,6 +204,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.InputLocation
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.IntoPathDescriptor;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.OutputColumn;
import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.SeriesScanOptions;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ForecastModelInferenceDescriptor;
import org.apache.iotdb.db.queryengine.plan.statement.component.FillPolicy;
import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByKey;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
@@ -214,6 +217,8 @@ import org.apache.iotdb.db.queryengine.transformation.dag.udf.UDTFContext;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.read.TimeValuePair;
import org.apache.iotdb.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.iotdb.tsfile.read.common.block.column.DoubleColumn;
+import org.apache.iotdb.tsfile.read.common.block.column.TimeColumn;
import org.apache.iotdb.tsfile.read.filter.basic.Filter;
import org.apache.iotdb.tsfile.read.filter.operator.Gt;
import org.apache.iotdb.tsfile.read.filter.operator.GtEq;
@@ -243,6 +248,7 @@ import static com.google.common.base.Preconditions.checkArgument;
import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.isSameNode;
import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.calculateMaxAggregationResultSize;
import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.calculateMaxAggregationResultSizeForLastQuery;
+import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.getOutputColumnSizePerLine;
import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.initTimeRangeIterator;
import static org.apache.iotdb.db.queryengine.plan.expression.leaf.TimestampOperand.TIMESTAMP_EXPRESSION_STRING;
import static org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.SeriesScanOptions.updateFilterUsingTTL;
@@ -1623,6 +1629,49 @@ public class OperatorTreeGenerator extends PlanVisitor<Operator, LocalExecutionP
MergeSortComparator.getComparator(sortItemList, sortItemIndexList, sortItemDataTypeList));
}
+ @Override
+ public Operator visitForecast(ForecastNode node, LocalExecutionPlanContext context) {
+ Operator child = node.getChild().accept(this, context);
+ OperatorContext operatorContext =
+ context
+ .getDriverContext()
+ .addOperatorContext(
+ context.getNextOperatorId(),
+ node.getPlanNodeId(),
+ ForecastOperator.class.getSimpleName());
+
+ ForecastModelInferenceDescriptor forecastModelInferenceDescriptor =
+ node.getModelInferenceDescriptor();
+
+ List<TSDataType> inputTypeList = forecastModelInferenceDescriptor.getInputTypeList();
+ int modelInputLength = forecastModelInferenceDescriptor.getModelInputLength();
+ long timeValueColumnsSizePerLine = TimeColumn.SIZE_IN_BYTES_PER_POSITION;
+ for (TSDataType dataType : inputTypeList) {
+ timeValueColumnsSizePerLine += getOutputColumnSizePerLine(dataType);
+ }
+ long maxRetainedSize = timeValueColumnsSizePerLine * modelInputLength;
+
+ int expectedPredictLength = forecastModelInferenceDescriptor.getExpectedPredictLength();
+ int outputColumnNum = forecastModelInferenceDescriptor.getPredictIndexList().size();
+ long maxReturnSize =
+ (TimeColumn.SIZE_IN_BYTES_PER_POSITION
+ + (long) outputColumnNum * DoubleColumn.SIZE_IN_BYTES_PER_POSITION)
+ * expectedPredictLength;
+
+ context.getTimeSliceAllocator().recordExecutionWeight(operatorContext, 1);
+
+ return new ForecastOperator(
+ operatorContext,
+ child,
+ forecastModelInferenceDescriptor.getModelPath(),
+ inputTypeList,
+ node.getChild().getOutputColumnNames(),
+ expectedPredictLength,
+ FragmentInstanceManager.getInstance().getModelInferenceExecutor(),
+ maxRetainedSize,
+ maxReturnSize);
+ }
+
@Override
public Operator visitInto(IntoNode node, LocalExecutionPlanContext context) {
Operator child = node.getChild().accept(this, context);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanGraphPrinter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanGraphPrinter.java
index ff72ddf559c..77971989322 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanGraphPrinter.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanGraphPrinter.java
@@ -45,6 +45,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TransformN
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryCollectNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml.ForecastNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.ShuffleSinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedLastQueryScanNode;
@@ -451,6 +452,17 @@ public class PlanGraphPrinter extends PlanVisitor<List<String>, PlanGraphPrinter
return render(node, boxValue, context);
}
+ @Override
+ public List<String> visitForecast(ForecastNode node, GraphContext context) {
+ List<String> boxValue = new ArrayList<>();
+ boxValue.add(String.format("Forecast-%s", node.getPlanNodeId().getId()));
+ boxValue.add("Output: ");
+ for (String outputColumnName : node.getOutputColumnNames()) {
+ boxValue.add(String.format(" %s", outputColumnName));
+ }
+ return render(node, boxValue, context);
+ }
+
private String printRegion(TRegionReplicaSet regionReplicaSet) {
return String.format(
"Partition: %s",
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java
index 2364aef9d7e..180775622a8 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java
@@ -78,6 +78,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TransformN
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryCollectNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml.ForecastNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.ShuffleSinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedLastQueryScanNode;
@@ -180,7 +181,7 @@ public enum PlanNodeType {
LOGICAL_VIEW_SCHEMA_SCAN((short) 77),
ALTER_LOGICAL_VIEW((short) 78),
PIPE_ENRICHED_INSERT((short) 79),
- ;
+ FORECAST((short) 80);
public static final int BYTES = Short.BYTES;
@@ -385,6 +386,8 @@ public enum PlanNodeType {
return AlterLogicalViewNode.deserialize(buffer);
case 79:
return PipeEnrichedInsertNode.deserialize(buffer);
+ case 80:
+ return ForecastNode.deserialize(buffer);
default:
throw new IllegalArgumentException("Invalid node type: " + nodeType);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java
index e6eb22fc5a2..f1b9943936b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java
@@ -78,6 +78,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.TransformN
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryCollectNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.last.LastQueryNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml.ForecastNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.ShuffleSinkNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.AlignedLastQueryScanNode;
@@ -192,6 +193,10 @@ public abstract class PlanVisitor<R, C> {
return visitSingleChildProcess(node, context);
}
+ public R visitForecast(ForecastNode node, C context) {
+ return visitSingleChildProcess(node, context);
+ }
+
// multi child --------------------------------------------------------------------------------
public R visitMultiChildProcess(MultiChildProcessNode node, C context) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/ml/ForecastNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/ml/ForecastNode.java
new file mode 100644
index 00000000000..9edccfee589
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/ml/ForecastNode.java
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ml;
+
+import org.apache.iotdb.db.queryengine.plan.expression.Expression;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeType;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.SingleChildProcessNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ForecastModelInferenceDescriptor;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+public class ForecastNode extends SingleChildProcessNode {
+
+ private final ForecastModelInferenceDescriptor modelInferenceDescriptor;
+
+ private List<String> outputColumnNames;
+
+ public ForecastNode(
+ PlanNodeId id, PlanNode child, ForecastModelInferenceDescriptor modelInferenceDescriptor) {
+ super(id, child);
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ }
+
+ public ForecastNode(PlanNodeId id, ForecastModelInferenceDescriptor modelInferenceDescriptor) {
+ super(id);
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ }
+
+ public ForecastModelInferenceDescriptor getModelInferenceDescriptor() {
+ return modelInferenceDescriptor;
+ }
+
+ @Override
+ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
+ return visitor.visitForecast(this, context);
+ }
+
+ @Override
+ public PlanNode clone() {
+ return new ForecastNode(getPlanNodeId(), child, modelInferenceDescriptor);
+ }
+
+ @Override
+ public List<String> getOutputColumnNames() {
+ if (outputColumnNames == null) {
+ outputColumnNames = new ArrayList<>();
+ for (Expression expression : modelInferenceDescriptor.getModelInferenceOutputExpressions()) {
+ outputColumnNames.add(expression.toString());
+ }
+ }
+ return outputColumnNames;
+ }
+
+ @Override
+ protected void serializeAttributes(ByteBuffer byteBuffer) {
+ PlanNodeType.FORECAST.serialize(byteBuffer);
+ modelInferenceDescriptor.serialize(byteBuffer);
+ }
+
+ @Override
+ protected void serializeAttributes(DataOutputStream stream) throws IOException {
+ PlanNodeType.FORECAST.serialize(stream);
+ modelInferenceDescriptor.serialize(stream);
+ }
+
+ public static ForecastNode deserialize(ByteBuffer buffer) {
+ ForecastModelInferenceDescriptor modelInferenceDescriptor =
+ ForecastModelInferenceDescriptor.deserialize(buffer);
+ PlanNodeId planNodeId = PlanNodeId.deserialize(buffer);
+ return new ForecastNode(planNodeId, modelInferenceDescriptor);
+ }
+
+ @Override
+ public String toString() {
+ return "ForecastNode-" + this.getPlanNodeId();
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ if (!super.equals(o)) {
+ return false;
+ }
+ ForecastNode that = (ForecastNode) o;
+ return modelInferenceDescriptor.equals(that.modelInferenceDescriptor);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), modelInferenceDescriptor);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ForecastModelInferenceDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ForecastModelInferenceDescriptor.java
new file mode 100644
index 00000000000..f7b075d14c6
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ForecastModelInferenceDescriptor.java
@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model;
+
+import org.apache.iotdb.commons.model.ForecastModeInformation;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
+import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Objects;
+
+import static org.apache.iotdb.db.utils.constant.SqlConstant.MODEL_ID;
+import static org.apache.iotdb.db.utils.constant.SqlConstant.PREDICT_LENGTH;
+
+public class ForecastModelInferenceDescriptor extends ModelInferenceDescriptor {
+
+ private final List<TSDataType> inputTypeList;
+ private final List<Integer> predictIndexList;
+
+ private final int modelInputLength;
+ private final int modelPredictLength;
+ private int expectedPredictLength;
+
+ private LinkedHashMap<String, String> outputAttributes;
+
+ public ForecastModelInferenceDescriptor(
+ ModelInferenceFunction functionType, ForecastModeInformation modelInformation) {
+ super(functionType, modelInformation);
+ this.inputTypeList = modelInformation.getInputTypeList();
+ this.predictIndexList = modelInformation.getPredictIndexList();
+ this.modelInputLength = modelInformation.getInputLength();
+ this.expectedPredictLength = this.modelPredictLength = modelInformation.getPredictLength();
+ }
+
+ public ForecastModelInferenceDescriptor(ByteBuffer buffer) {
+ super(buffer);
+ int listSize = ReadWriteIOUtils.readInt(buffer);
+ this.inputTypeList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ this.inputTypeList.add(TSDataType.deserializeFrom(buffer));
+ }
+ listSize = ReadWriteIOUtils.readInt(buffer);
+ this.predictIndexList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ this.predictIndexList.add(ReadWriteIOUtils.readInt(buffer));
+ }
+ this.modelInputLength = ReadWriteIOUtils.readInt(buffer);
+ this.modelPredictLength = ReadWriteIOUtils.readInt(buffer);
+ this.expectedPredictLength = ReadWriteIOUtils.readInt(buffer);
+ }
+
+ public List<Integer> getPredictIndexList() {
+ return predictIndexList;
+ }
+
+ public List<TSDataType> getInputTypeList() {
+ return inputTypeList;
+ }
+
+ public int getModelInputLength() {
+ return modelInputLength;
+ }
+
+ public int getModelPredictLength() {
+ return modelPredictLength;
+ }
+
+ public int getExpectedPredictLength() {
+ return expectedPredictLength;
+ }
+
+ public void setExpectedPredictLength(int expectedPredictLength) {
+ this.expectedPredictLength = expectedPredictLength;
+ }
+
+ @Override
+ public LinkedHashMap<String, String> getOutputAttributes() {
+ if (outputAttributes == null) {
+ outputAttributes = new LinkedHashMap<>();
+ outputAttributes.put(MODEL_ID, modelId);
+ if (expectedPredictLength != modelPredictLength) {
+ outputAttributes.put(PREDICT_LENGTH, String.valueOf(expectedPredictLength));
+ }
+ }
+ return outputAttributes;
+ }
+
+ @Override
+ public void serialize(ByteBuffer byteBuffer) {
+ super.serialize(byteBuffer);
+ ReadWriteIOUtils.write(inputTypeList.size(), byteBuffer);
+ for (TSDataType dataType : inputTypeList) {
+ dataType.serializeTo(byteBuffer);
+ }
+ ReadWriteIOUtils.write(predictIndexList.size(), byteBuffer);
+ for (Integer index : predictIndexList) {
+ ReadWriteIOUtils.write(index, byteBuffer);
+ }
+ ReadWriteIOUtils.write(modelInputLength, byteBuffer);
+ ReadWriteIOUtils.write(modelPredictLength, byteBuffer);
+ ReadWriteIOUtils.write(expectedPredictLength, byteBuffer);
+ }
+
+ @Override
+ public void serialize(DataOutputStream stream) throws IOException {
+ super.serialize(stream);
+ ReadWriteIOUtils.write(inputTypeList.size(), stream);
+ for (TSDataType dataType : inputTypeList) {
+ dataType.serializeTo(stream);
+ }
+ ReadWriteIOUtils.write(predictIndexList.size(), stream);
+ for (Integer index : predictIndexList) {
+ ReadWriteIOUtils.write(index, stream);
+ }
+ ReadWriteIOUtils.write(modelInputLength, stream);
+ ReadWriteIOUtils.write(modelPredictLength, stream);
+ ReadWriteIOUtils.write(expectedPredictLength, stream);
+ }
+
+ public static ForecastModelInferenceDescriptor deserialize(ByteBuffer buffer) {
+ return new ForecastModelInferenceDescriptor(buffer);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ if (!super.equals(o)) {
+ return false;
+ }
+ ForecastModelInferenceDescriptor that = (ForecastModelInferenceDescriptor) o;
+ return modelInputLength == that.modelInputLength
+ && modelPredictLength == that.modelPredictLength
+ && expectedPredictLength == that.expectedPredictLength
+ && inputTypeList.equals(that.inputTypeList)
+ && predictIndexList.equals(that.predictIndexList);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ super.hashCode(),
+ inputTypeList,
+ predictIndexList,
+ modelInputLength,
+ modelPredictLength,
+ expectedPredictLength);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ModelInferenceDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ModelInferenceDescriptor.java
new file mode 100644
index 00000000000..dfbb12545e3
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ModelInferenceDescriptor.java
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model;
+
+import org.apache.iotdb.commons.model.ModelInformation;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
+import org.apache.iotdb.db.queryengine.plan.expression.multi.FunctionExpression;
+import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Objects;
+
+public abstract class ModelInferenceDescriptor {
+
+ protected final ModelInferenceFunction functionType;
+
+ protected final String modelId;
+
+ protected final String modelPath;
+
+ protected List<FunctionExpression> modelInferenceOutputExpressions;
+
+ public ModelInferenceDescriptor(
+ ModelInferenceFunction functionType, ModelInformation modelInformation) {
+ this.functionType = functionType;
+ this.modelId = modelInformation.getModelId();
+ this.modelPath = modelInformation.getModelPath();
+ }
+
+ public ModelInferenceDescriptor(ByteBuffer buffer) {
+ this.functionType = ModelInferenceFunction.values()[ReadWriteIOUtils.readInt(buffer)];
+ this.modelId = ReadWriteIOUtils.readString(buffer);
+ this.modelPath = ReadWriteIOUtils.readString(buffer);
+ }
+
+ public ModelInferenceFunction getFunctionType() {
+ return functionType;
+ }
+
+ public String getModelId() {
+ return modelId;
+ }
+
+ public String getModelPath() {
+ return modelPath;
+ }
+
+ public List<FunctionExpression> getModelInferenceOutputExpressions() {
+ return modelInferenceOutputExpressions;
+ }
+
+ public void setModelInferenceOutputExpressions(
+ List<FunctionExpression> modelInferenceOutputExpressions) {
+ this.modelInferenceOutputExpressions = modelInferenceOutputExpressions;
+ }
+
+ public abstract LinkedHashMap<String, String> getOutputAttributes();
+
+ public void serialize(ByteBuffer byteBuffer) {
+ ReadWriteIOUtils.write(functionType.ordinal(), byteBuffer);
+ ReadWriteIOUtils.write(modelId, byteBuffer);
+ ReadWriteIOUtils.write(modelPath, byteBuffer);
+ }
+
+ public void serialize(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(functionType.ordinal(), stream);
+ ReadWriteIOUtils.write(modelId, stream);
+ ReadWriteIOUtils.write(modelPath, stream);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ ModelInferenceDescriptor that = (ModelInferenceDescriptor) o;
+ return functionType == that.functionType
+ && modelId.equals(that.modelId)
+ && modelPath.equals(that.modelPath)
+ && modelInferenceOutputExpressions.equals(that.modelInferenceOutputExpressions);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(functionType, modelId, modelPath, modelInferenceOutputExpressions);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java
index 17bee5aa40b..eefeca85510 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java
@@ -79,7 +79,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.UnSetTTLStatement
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.DropModelStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
-import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrailsStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowTrialsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.DropPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.ShowPipesStatement;
@@ -280,8 +280,8 @@ public abstract class StatementVisitor<R, C> {
return visitStatement(showModelsModelStatement, context);
}
- public R visitShowTrails(ShowTrailsStatement showTrailsStatement, C context) {
- return visitStatement(showTrailsStatement, context);
+ public R visitShowTrials(ShowTrialsStatement showTrialsStatement, C context) {
+ return visitStatement(showTrialsStatement, context);
}
/** Data Manipulation Language (DML) */
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/ResultColumn.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/ResultColumn.java
index e8db3f8a5ee..e7299e60e63 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/ResultColumn.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/ResultColumn.java
@@ -77,7 +77,7 @@ public class ResultColumn extends StatementNode {
public ResultColumn(Expression expression, ColumnType columnType) {
this.expression = expression;
this.columnType = columnType;
- alias = null;
+ this.alias = null;
}
public Expression getExpression() {
@@ -127,6 +127,7 @@ public class ResultColumn extends StatementNode {
public enum ColumnType {
RAW,
AGGREGATION,
- CONSTANT
+ CONSTANT,
+ MODEL_INFERENCE
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/SelectComponent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/SelectComponent.java
index 76c5d6d47ee..dfe3a6d65cf 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/SelectComponent.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/SelectComponent.java
@@ -36,6 +36,8 @@ public class SelectComponent extends StatementNode {
private boolean hasBuiltInAggregationFunction = false;
+ private boolean hasModelInferenceFunction = false;
+
protected List<ResultColumn> resultColumns = new ArrayList<>();
private Map<String, Expression> aliasToColumnMap;
@@ -48,15 +50,21 @@ public class SelectComponent extends StatementNode {
return zoneId;
}
- public boolean isHasBuiltInAggregationFunction() {
+ public boolean hasBuiltInAggregationFunction() {
return hasBuiltInAggregationFunction;
}
+ public boolean hasModelInferenceFunction() {
+ return hasModelInferenceFunction;
+ }
+
public void addResultColumn(ResultColumn resultColumn) {
resultColumns.add(resultColumn);
ResultColumn.ColumnType columnType = resultColumn.getColumnType();
if (columnType == ResultColumn.ColumnType.AGGREGATION) {
hasBuiltInAggregationFunction = true;
+ } else if (columnType == ResultColumn.ColumnType.MODEL_INFERENCE) {
+ hasModelInferenceFunction = true;
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/QueryStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/QueryStatement.java
index 06458367fd4..0ea683d4de1 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/QueryStatement.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/QueryStatement.java
@@ -53,6 +53,7 @@ import java.util.List;
import java.util.Set;
import static org.apache.iotdb.db.utils.constant.SqlConstant.COUNT_TIME;
+import static org.apache.iotdb.db.utils.constant.SqlConstant.MODEL_ID;
/**
* Base class of SELECT statement.
@@ -294,7 +295,7 @@ public class QueryStatement extends Statement {
}
public boolean isAggregationQuery() {
- return selectComponent.isHasBuiltInAggregationFunction();
+ return selectComponent.hasBuiltInAggregationFunction();
}
public boolean isGroupByLevel() {
@@ -513,6 +514,10 @@ public class QueryStatement extends Statement {
public static final String RAW_AGGREGATION_HYBRID_QUERY_ERROR_MSG =
"Raw data and aggregation hybrid query is not supported.";
+ public boolean isModelInferenceQuery() {
+ return selectComponent.hasModelInferenceFunction();
+ }
+
public static final String COUNT_TIME_NOT_SUPPORT_GROUP_BY_LEVEL =
"Count_time aggregation function using with group by level is not supported.";
@@ -531,6 +536,51 @@ public class QueryStatement extends Statement {
@SuppressWarnings({"squid:S3776", "squid:S6541"}) // Suppress high Cognitive Complexity warning
public void semanticCheck() {
+ if (isModelInferenceQuery()) {
+ if (selectComponent.getResultColumns().size() > 1) {
+ throw new SemanticException("Model inference query can't be used with other queries now.");
+ }
+
+ Expression modelInferenceExpression =
+ selectComponent.getResultColumns().get(0).getExpression();
+ if (!(modelInferenceExpression instanceof FunctionExpression
+ && ((FunctionExpression) modelInferenceExpression).isModelInferenceFunction())) {
+ throw new SemanticException("Model inference only supports forecast function now.");
+ }
+ if (!((FunctionExpression) modelInferenceExpression)
+ .getFunctionAttributes()
+ .containsKey(MODEL_ID)) {
+ throw new SemanticException("Model inference function must indicate the model id.");
+ }
+ if (!ExpressionAnalyzer.searchAggregationExpressions(modelInferenceExpression).isEmpty()) {
+ throw new SemanticException(
+ "Expression "
+ + modelInferenceExpression.getExpressionString()
+ + " doesn't refer to any data.");
+ }
+
+ if (hasHaving()
+ || isGroupBy()
+ || isGroupByLevel()
+ || isGroupByTag()
+ || isAlignByDevice()
+ || isLastQuery()
+ || seriesLimit > 0
+ || seriesOffset > 0
+ || isSelectInto()
+ || isOrderByDevice()
+ || isOrderByTimeseries()) {
+ throw new SemanticException("Model inference query can't be used with other queries now.");
+ }
+
+ if (orderByComponent != null
+ && (!orderByComponent.isOrderByTime()
+ || orderByComponent.getTimeOrder() != Ordering.ASC)) {
+ throw new SemanticException(
+ "The order by clause of model inference query must be ORDER BY TIME ASC.");
+ }
+ }
+
if (isAggregationQuery()) {
if (groupByComponent != null && isGroupByLevel()) {
throw new SemanticException("GROUP BY CLAUSES doesn't support GROUP BY LEVEL now.");
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/CreateModelStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/CreateModelStatement.java
index 5065e367fff..b7e50ee72e2 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/CreateModelStatement.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/CreateModelStatement.java
@@ -19,7 +19,7 @@
package org.apache.iotdb.db.queryengine.plan.statement.metadata.model;
-import org.apache.iotdb.common.rpc.thrift.ModelTask;
+import org.apache.iotdb.common.rpc.thrift.TaskType;
import org.apache.iotdb.commons.path.PartialPath;
import org.apache.iotdb.db.exception.sql.SemanticException;
import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
@@ -28,16 +28,25 @@ import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import static org.apache.iotdb.commons.model.ModelInformation.MODEL_TYPE;
+import static org.apache.iotdb.commons.model.ModelInformation.TASK_TYPE;
+
public class CreateModelStatement extends Statement implements IConfigStatement {
private String modelId;
- private boolean isAuto;
- private Map<String, String> attributes;
- private QueryStatement queryStatement;
+
+ private Map<String, String> options;
+
+ private Map<String, String> hyperparameters;
+
+ private QueryStatement datasetStatement;
+
+ private final List<String> requiredOptions = Arrays.asList(TASK_TYPE, MODEL_TYPE);
public CreateModelStatement() {
// do nothing
@@ -51,44 +60,54 @@ public class CreateModelStatement extends Statement implements IConfigStatement
this.modelId = modelId;
}
- public boolean isAuto() {
- return isAuto;
+ public Map<String, String> getOptions() {
+ return options;
}
- public void setAuto(boolean auto) {
- isAuto = auto;
+ public void setOptions(Map<String, String> options) {
+ this.options = options;
}
- public Map<String, String> getAttributes() {
- return attributes;
+ public Map<String, String> getHyperparameters() {
+ return hyperparameters;
}
- public void setAttributes(Map<String, String> attributes) {
- this.attributes = attributes;
+ public void setHyperparameters(Map<String, String> hyperparameters) {
+ this.hyperparameters = hyperparameters;
}
- public QueryStatement getQueryStatement() {
- return queryStatement;
+ public QueryStatement getDatasetStatement() {
+ return datasetStatement;
}
- public void setQueryStatement(QueryStatement queryStatement) {
- this.queryStatement = queryStatement;
+ public void setDatasetStatement(QueryStatement queryBody) {
+ this.datasetStatement = queryBody;
}
- public ModelTask getModelTask() {
- return ModelTask.valueOf(attributes.get("model_task").toUpperCase());
- }
-
- public String getModelType() {
- return attributes.get("model_type");
+ public TaskType getTaskType() {
+ try {
+ return TaskType.valueOf(options.get("task_type").toUpperCase());
+ } catch (IllegalArgumentException e) {
+ throw new SemanticException("Unknown task type: " + options.get("task_type"));
+ }
}
public void semanticCheck() {
- if (!attributes.containsKey("model_task")) {
- throw new SemanticException("The attribute `model_task` must be specified.");
+ for (String requiredOption : requiredOptions) {
+ if (!options.containsKey(requiredOption)) {
+ throw new SemanticException("The option `" + requiredOption + "` must be specified.");
+ }
+ }
+ if (datasetStatement.isAlignByDevice()) {
+ throw new SemanticException(
+ "Currently the sql statement in CREATE MODEL does not support align by device.");
+ }
+ if (datasetStatement.isLastQuery()) {
+ throw new SemanticException("The sql statement in CREATE MODEL does not support LAST query.");
}
- if (!attributes.containsKey("model_type")) {
- throw new SemanticException("The attribute `model_type` must be specified.");
+ if (datasetStatement.isAggregationQuery() && !datasetStatement.isGroupByTime()) {
+ throw new SemanticException(
+ "The aggregation sql in CREATE MODEL only supports GROUP BY TIME now.");
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrailsStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrialsStatement.java
similarity index 91%
rename from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrailsStatement.java
rename to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrialsStatement.java
index f590047ba14..00c60a2fe5f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrailsStatement.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowTrialsStatement.java
@@ -28,11 +28,11 @@ import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
import java.util.Collections;
import java.util.List;
-public class ShowTrailsStatement extends Statement implements IConfigStatement {
+public class ShowTrialsStatement extends Statement implements IConfigStatement {
private final String modelId;
- public ShowTrailsStatement(String modelId) {
+ public ShowTrialsStatement(String modelId) {
this.modelId = modelId;
}
@@ -52,6 +52,6 @@ public class ShowTrailsStatement extends Statement implements IConfigStatement {
@Override
public <R, C> R accept(StatementVisitor<R, C> visitor, C context) {
- return visitor.visitShowTrails(this, context);
+ return visitor.visitShowTrials(this, context);
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/MLNodeRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/MLNodeRPCService.java
index 018d98ecd23..959d55c8f48 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/MLNodeRPCService.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/MLNodeRPCService.java
@@ -30,8 +30,6 @@ import org.apache.iotdb.db.protocol.thrift.handler.MLNodeRPCServiceThriftHandler
import org.apache.iotdb.db.protocol.thrift.impl.MLNodeRPCServiceImpl;
import org.apache.iotdb.mpp.rpc.thrift.IMLNodeInternalRPCService;
-import java.lang.reflect.InvocationTargetException;
-
public class MLNodeRPCService extends ThriftService implements MLNodeRPCServiceMBean {
private MLNodeRPCServiceImpl impl;
@@ -44,9 +42,7 @@ public class MLNodeRPCService extends ThriftService implements MLNodeRPCServiceM
}
@Override
- public void initTProcessor()
- throws ClassNotFoundException, IllegalAccessException, InstantiationException,
- NoSuchMethodException, InvocationTargetException {
+ public void initTProcessor() {
impl = new MLNodeRPCServiceImpl();
initSyncedServiceImpl(null);
processor = new IMLNodeInternalRPCService.Processor<>(impl);
@@ -67,13 +63,11 @@ public class MLNodeRPCService extends ThriftService implements MLNodeRPCServiceM
config.getRpcMaxConcurrentClientNum(),
config.getThriftServerAwaitTimeForStopService(),
new MLNodeRPCServiceThriftHandler(impl),
- // TODO: hard coded compress strategy
false);
} catch (RPCServiceException e) {
throw new IllegalAccessException(e.getMessage());
}
thriftServiceThread.setName(ThreadName.MLNODE_RPC_SERVICE.getName());
- // TODO: metricService
}
@Override
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/exception/ModelInferenceProcessException.java
similarity index 80%
copy from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
copy to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/exception/ModelInferenceProcessException.java
index 304d209817f..1ddb212f113 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/exception/ModelInferenceProcessException.java
@@ -17,11 +17,11 @@
* under the License.
*/
-package org.apache.iotdb.confignode.procedure.state.model;
+package org.apache.iotdb.db.exception;
-public enum CreateModelState {
- INIT,
- VALIDATED,
- CONFIG_NODE_ACTIVE,
- ML_NODE_ACTIVE
+public class ModelInferenceProcessException extends RuntimeException {
+
+ public ModelInferenceProcessException(String message) {
+ super(message);
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/TimestampPrecisionUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/TimestampPrecisionUtils.java
index e613dacf908..ee1a61dac18 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/TimestampPrecisionUtils.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/TimestampPrecisionUtils.java
@@ -23,7 +23,7 @@ import org.apache.iotdb.commons.conf.CommonDescriptor;
import java.util.concurrent.TimeUnit;
public class TimestampPrecisionUtils {
- private static final String timestampPrecision =
+ private static final String TIMESTAMP_PRECISION =
CommonDescriptor.getInstance().getConfig().getTimestampPrecision();
@FunctionalInterface
@@ -34,16 +34,16 @@ public class TimestampPrecisionUtils {
private static final ConvertFunction<Long, TimeUnit, Long> convertFunction;
static {
- switch (timestampPrecision) {
+ switch (TIMESTAMP_PRECISION) {
case "ms":
convertFunction = TimeUnit.MILLISECONDS::convert;
break;
- case "ns":
- convertFunction = TimeUnit.NANOSECONDS::convert;
- break;
case "us":
convertFunction = TimeUnit.MICROSECONDS::convert;
break;
+ case "ns":
+ convertFunction = TimeUnit.NANOSECONDS::convert;
+ break;
// this case will never reach
default:
convertFunction = null;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/constant/SqlConstant.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/constant/SqlConstant.java
index afd47cc8c51..e98b34c7f44 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/constant/SqlConstant.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/constant/SqlConstant.java
@@ -81,6 +81,9 @@ public class SqlConstant {
public static final String SUBSTRING_IS_STANDARD = "isStandard";
public static final String SUBSTRING_FOR = "FOR";
+ public static final String MODEL_ID = "model_id";
+ public static final String PREDICT_LENGTH = "predict_length";
+
public static String[] getSingleRootArray() {
return SINGLE_ROOT_ARRAY;
}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
index ff0ec603365..17dc92f3adf 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
@@ -403,11 +403,10 @@ public class StatementGeneratorTest {
}
@Test
- public void testFetchTimeseries() throws IllegalPathException {
- TFetchTimeseriesReq req = new TFetchTimeseriesReq(Collections.singletonList("root.sg.d1.s1"));
- QueryStatement statement =
- (QueryStatement) StatementGenerator.createStatement(req, ZonedDateTime.now().getOffset());
- assertEquals(Collections.singletonList(new PartialPath("root.sg.d1.s1")), statement.getPaths());
+ public void testFetchTimeseries() {
+ TFetchTimeseriesReq req = new TFetchTimeseriesReq("select * from root.sg.d1.s1");
+ Statement statement = StatementGenerator.createStatement(req, ZonedDateTime.now().getOffset());
+ assertEquals("root.sg.d1.s1.*", statement.getPaths().get(0).getFullPath());
}
@Test
diff --git a/iotdb-core/mlnode/.gitignore b/iotdb-core/mlnode/.gitignore
index 4a6f6eac650..cc99cc4f975 100644
--- a/iotdb-core/mlnode/.gitignore
+++ b/iotdb-core/mlnode/.gitignore
@@ -5,4 +5,5 @@
/dist/
# generated by MLNode
-*.pt
\ No newline at end of file
+*.pt
+/rebuild.bat
diff --git a/iotdb-core/mlnode/README.md b/iotdb-core/mlnode/README.md
index d892d2273e6..062db10f7bf 100644
--- a/iotdb-core/mlnode/README.md
+++ b/iotdb-core/mlnode/README.md
@@ -26,11 +26,11 @@ For developers, you can start an ML Node through the following steps.
- Step 1: build project
```shell
-mvn clean package -DskipUTs -pl mlnode -am
+mvn clean package -DskipUTs -pl iotdb-core/mlnode -am
```
```shell
-cd mlnode
+cd iotdb-core/mlnode
poetry build
```
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/factory.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/factory.py
new file mode 100644
index 00000000000..fb25a27cac9
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/factory.py
@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import Dict
+
+import torch.nn as nn
+
+from iotdb.mlnode.algorithm.hyperparameter import HyperparameterName
+from iotdb.mlnode.algorithm.models.forecast.dlinear import (DLinear,
+ DLinearIndividual)
+from iotdb.mlnode.algorithm.models.forecast.nbeats import NBeats
+from iotdb.mlnode.constant import ForecastModelType, OptionsKey
+from iotdb.mlnode.exception import BadConfigValueError
+from iotdb.mlnode.parser import ForecastTaskOptions
+
+
+def create_forecast_model(
+ model_type: ForecastModelType,
+ model_configs: Dict,
+) -> nn.Module:
+ """
+ Factory method for all support forecasting models
+ the given arguments is common configs shared by all forecasting models
+ """
+ if model_type.value not in ForecastModelType.values():
+ raise BadConfigValueError('model_name', f'It should be one of {ForecastModelType.values()}')
+
+ if model_type == ForecastModelType.DLINEAR:
+ return DLinear(kernel_size=model_configs[HyperparameterName.KERNEL_SIZE.value],
+ input_len=model_configs[OptionsKey.INPUT_LENGTH.value],
+ pred_len=model_configs[OptionsKey.PREDICT_LENGTH.value],
+ input_vars=model_configs[HyperparameterName.INPUT_VARS.value])
+ elif model_type == ForecastModelType.DLINEAR_INDIVIDUAL:
+ return DLinearIndividual(kernel_size=model_configs[HyperparameterName.KERNEL_SIZE.value],
+ input_len=model_configs[OptionsKey.INPUT_LENGTH.value],
+ pred_len=model_configs[OptionsKey.PREDICT_LENGTH.value],
+ input_vars=model_configs[HyperparameterName.INPUT_VARS.value]
+ )
+ elif model_type == ForecastModelType.NBEATS:
+ return NBeats(
+ block_type=model_configs[HyperparameterName.BLOCK_TYPE.value],
+ d_model=model_configs[HyperparameterName.D_MODEL.value],
+ inner_layers=model_configs[HyperparameterName.INNER_LAYERS.value],
+ outer_layers=model_configs[HyperparameterName.OUTER_LAYERS.value],
+ input_len=model_configs[OptionsKey.INPUT_LENGTH.value],
+ pred_len=model_configs[OptionsKey.PREDICT_LENGTH.value],
+ input_vars=model_configs[HyperparameterName.INPUT_VARS.value]
+ )
+ else:
+ raise BadConfigValueError('model_name', model_type,
+ f'It should be one of {ForecastModelType.values()}')
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/hyperparameter.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/hyperparameter.py
new file mode 100644
index 00000000000..44b8feb9f10
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/hyperparameter.py
@@ -0,0 +1,407 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from abc import abstractmethod
+from typing import Dict, List, Optional, Tuple
+
+import optuna
+
+from iotdb.mlnode.algorithm.validator import NumberRangeValidator, Validator
+from iotdb.mlnode.constant import HyperparameterName, ForecastModelType
+from iotdb.mlnode.exception import BadConfigValueError, UnsupportedError
+from iotdb.mlnode.parser import ForecastTaskOptions, TaskOptions
+
+
+class Hyperparameter(object):
+
+ def __init__(self, name: str):
+ """
+ Args:
+ name: name of the hyperparameter
+ """
+ self._name = name
+
+ @abstractmethod
+ def get_default_value(self):
+ raise NotImplementedError
+
+ @abstractmethod
+ def validate_value(self, value):
+ raise NotImplementedError
+
+ @abstractmethod
+ def validate_range(self, min_value, max_value):
+ raise NotImplementedError
+
+ @abstractmethod
+ def suggest_parameter(self, optuna_suggest: optuna.Trial):
+ raise NotImplementedError
+
+ @abstractmethod
+ def parse(self, string_value: str):
+ raise NotImplementedError
+
+
+class IntHyperparameter(Hyperparameter):
+ def __init__(self, name: str,
+ log: bool,
+ default_value: int,
+ value_validators: List[Validator],
+ default_low: int,
+ low_validators: List[Validator],
+ default_high: int,
+ high_validators: List[Validator],
+ step: Optional[int] = None,
+ tuning: bool = False
+ ):
+ super(IntHyperparameter, self).__init__(name)
+ self.__log = log
+ self.__default_value = default_value
+ self.__value_validators = value_validators
+ self.__default_low = default_low
+ self.__low_validators = low_validators
+ self.__default_high = default_high
+ self.__high_validators = high_validators
+ self.__step = step
+ self.__tuning = tuning
+
+ def get_default_value(self):
+ return self.__default_value
+
+ def __validate(self, value, validators: List[Validator] = None):
+ try:
+ for validator in validators:
+ validator.validate(int(value))
+ except Exception as e:
+ raise BadConfigValueError(self._name, value, str(e))
+
+ def validate_value(self, value):
+ self.__validate(value, self.__value_validators)
+
+ def validate_range(self, min_value: int, max_value: int):
+ self.__validate(min_value, self.__low_validators)
+ self.__validate(max_value, self.__high_validators)
+
+ def suggest_parameter(self, optuna_suggest: optuna.Trial):
+ if not self.__tuning:
+ return self.__default_value
+ return optuna_suggest.suggest_int(
+ name=self._name,
+ low=self.__default_low,
+ high=self.__default_high,
+ step=self.__step,
+ log=self.__log
+ )
+
+ def parse(self, string_value: str):
+ return int(string_value)
+
+
+class FloatHyperparameter(Hyperparameter):
+ def __init__(self, name: str,
+ log: bool,
+ default_value: float,
+ value_validators: List[Validator],
+ default_low: float,
+ low_validators: List[Validator],
+ default_high: float,
+ high_validators: List[Validator],
+ step: Optional[float] = None,
+ tuning: bool = False):
+ super(FloatHyperparameter, self).__init__(name)
+ self.__log = log
+ self.__default_value = default_value
+ self.__value_validators = value_validators
+ self.__default_low = default_low
+ self.__low_validators = low_validators
+ self.__default_high = default_high
+ self.__high_validators = high_validators
+ self.__step = step
+ self.__tuning = tuning
+
+ def get_default_value(self):
+ return self.__default_value
+
+ def __validate(self, value, validators: List[Validator] = None):
+ try:
+ for validator in validators:
+ validator.validate(float(value))
+ except Exception as e:
+ raise BadConfigValueError(self._name, value, str(e))
+
+ def validate_value(self, value):
+ self.__validate(value, self.__value_validators)
+
+ def validate_range(self, min_value: float, max_value: float):
+ self.__validate(min_value, self.__low_validators)
+ self.__validate(max_value, self.__high_validators)
+
+ def suggest_parameter(self, optuna_suggest: optuna.Trial):
+ if not self.__tuning:
+ return self.__default_value
+ return optuna_suggest.suggest_float(
+ name=self._name,
+ low=self.__default_low,
+ high=self.__default_high,
+ step=self.__step,
+ log=self.__log
+ )
+
+ def parse(self, string_value: str):
+ return float(string_value)
+
+
+class StringHyperparameter(Hyperparameter):
+ def __init__(self, name: str,
+ default_value: str,
+ value_validators: List[Validator],
+ value_list: List[str],
+ tuning: bool = False):
+ super(StringHyperparameter, self).__init__(name)
+ self.__default_value = default_value
+ self.__value_validators = value_validators
+ self.__value_list = value_list
+ self.__tuning = tuning
+
+ def get_default_value(self):
+ return self.__default_value
+
+ def __validate(self, value, validators: List[Validator] = None):
+ try:
+ for validator in validators:
+ validator.validate(float(value))
+ except Exception as e:
+ raise BadConfigValueError(self._name, value, str(e))
+
+ def validate_value(self, value):
+ self.__validate(value, self.__value_validators)
+
+ def validate_range(self, min_value: float, max_value: float):
+ raise UnsupportedError("validate range in string hyperparameter")
+
+ def suggest_parameter(self, optuna_suggest: optuna.Trial):
+ if not self.__tuning:
+ return self.__default_value
+ return optuna_suggest.suggest_categorical(
+ name=self._name,
+ choices=self.__value_list
+ )
+
+ def parse(self, string_value: str):
+ return string_value
+
+
+class BooleanHyperparameter(Hyperparameter):
+ def __init__(self, name: str,
+ default_value: bool,
+ tuning: bool = False):
+ super(BooleanHyperparameter, self).__init__(name)
+ self.__default_value = default_value
+ self.__tuning = tuning
+
+ def get_default_value(self):
+ return self.__default_value
+
+ def validate_value(self, value):
+ # no need to validate value in boolean hyperparameter
+ pass
+
+ def validate_range(self, min_value: float, max_value: float):
+ # no need to validate range in boolean hyperparameter
+ pass
+
+ def suggest_parameter(self, optuna_suggest: optuna.Trial):
+ if not self.__tuning:
+ return self.__default_value
+ return optuna_suggest.suggest_categorical(
+ self._name,
+ [True, False]
+ )
+
+ def parse(self, string_value: str):
+ return bool(string_value)
+
+
+training_hyperparameter_map = {
+ HyperparameterName.LEARNING_RATE.name(): FloatHyperparameter(
+ name=HyperparameterName.LEARNING_RATE.name(),
+ log=True,
+ default_value=1e-3,
+ value_validators=[NumberRangeValidator(1e-8, 10)],
+ default_low=1e-5,
+ low_validators=[],
+ default_high=1e-1,
+ high_validators=[],
+ tuning=True
+ ),
+ HyperparameterName.EPOCHS.name(): IntHyperparameter(
+ name=HyperparameterName.EPOCHS.name(),
+ log=True,
+ default_value=10,
+ value_validators=[NumberRangeValidator(1, 10000)],
+ default_low=1,
+ low_validators=[],
+ default_high=100,
+ high_validators=[],
+ tuning=False
+ ),
+ HyperparameterName.BATCH_SIZE.name(): IntHyperparameter(
+ name=HyperparameterName.BATCH_SIZE.name(),
+ log=True,
+ default_value=32,
+ value_validators=[NumberRangeValidator(1, 8196)],
+ default_low=1,
+ low_validators=[],
+ default_high=8196,
+ high_validators=[],
+ tuning=False
+ ),
+ HyperparameterName.USE_GPU.name(): BooleanHyperparameter(
+ name=HyperparameterName.USE_GPU.name(),
+ default_value=False,
+ tuning=False
+ ),
+ HyperparameterName.NUM_WORKERS.name(): IntHyperparameter(
+ name=HyperparameterName.NUM_WORKERS.name(),
+ log=False,
+ default_value=0,
+ value_validators=[NumberRangeValidator(0, 8)],
+ default_low=0,
+ low_validators=[],
+ default_high=8,
+ high_validators=[],
+ tuning=False
+ )
+}
+
+dlinear_structure_hyperparameter_map = {
+ HyperparameterName.KERNEL_SIZE.value: IntHyperparameter(
+ name=HyperparameterName.KERNEL_SIZE.name(),
+ log=True,
+ default_value=25,
+ value_validators=[NumberRangeValidator(1, 1e10)],
+ default_low=5,
+ low_validators=[],
+ default_high=50,
+ high_validators=[]
+ )
+}
+
+nbeats_structure_hyperparameter_map = {
+ HyperparameterName.BLOCK_TYPE.value: StringHyperparameter(
+ name=HyperparameterName.BLOCK_TYPE.name(),
+ default_value='generic',
+ value_validators=[],
+ value_list=['generic'],
+ tuning=False
+ ),
+ HyperparameterName.D_MODEL.value: IntHyperparameter(
+ name=HyperparameterName.D_MODEL.name(),
+ log=True,
+ default_value=512,
+ value_validators=[NumberRangeValidator(4, 8192)],
+ default_low=2,
+ low_validators=[],
+ default_high=2048,
+ high_validators=[]
+ ),
+ HyperparameterName.INNER_LAYERS.value: IntHyperparameter(
+ name=HyperparameterName.INNER_LAYERS.name(),
+ log=False,
+ default_value=4,
+ value_validators=[NumberRangeValidator(1, 128)],
+ default_low=1,
+ low_validators=[],
+ default_high=128,
+ high_validators=[]
+ ),
+ HyperparameterName.OUTER_LAYERS.value: IntHyperparameter(
+ name=HyperparameterName.OUTER_LAYERS.name(),
+ log=False,
+ default_value=4,
+ value_validators=[NumberRangeValidator(1, 128)],
+ default_low=1,
+ low_validators=[],
+ default_high=128,
+ high_validators=[])
+}
+
+
+def get_structure_hyperparameter_map(model_type: ForecastModelType) -> Dict[str, Hyperparameter]:
+ """
+ Different model may have different structure hyperparameters.
+ This method returns the structure hyperparameter map for a given model type.
+ """
+ if model_type == ForecastModelType.DLINEAR:
+ return dlinear_structure_hyperparameter_map
+ elif model_type == ForecastModelType.DLINEAR_INDIVIDUAL:
+ return dlinear_structure_hyperparameter_map
+ elif model_type == ForecastModelType.NBEATS:
+ return nbeats_structure_hyperparameter_map
+ else:
+ raise NotImplementedError(f"Model type {model_type} is not supported yet")
+
+
+def parse_fixed_hyperparameters(
+ task_options: ForecastTaskOptions,
+ input_hyperparameters: Dict[str, str]
+) -> Tuple[Dict, Dict]:
+ """
+ Parse the input hyperparameters into model hyperparameters and task hyperparameters.
+ If the input hyperparameters contains hyperparameters that are not defined in the model or task,
+ use default value.
+ """
+ structure_hyperparameter_map = get_structure_hyperparameter_map(task_options.model_type)
+ model_hyperparameters = parse_dict(input_hyperparameters, structure_hyperparameter_map)
+ task_hyperparameters = parse_dict(input_hyperparameters, training_hyperparameter_map)
+ return model_hyperparameters, task_hyperparameters
+
+
+def parse_dict(input_hyperparameters: Dict[str, str], hyperparameter_template_map: Dict[str, Hyperparameter]) -> Dict:
+ hyperparameter = {}
+ for hyperparameter_name in hyperparameter_template_map.keys():
+ hyperparameter_template = hyperparameter_template_map[hyperparameter_name]
+
+ # if user defines current hyperparameter
+ if hyperparameter_name in input_hyperparameters.keys():
+ value = hyperparameter_template.parse(input_hyperparameters[hyperparameter_name])
+ hyperparameter_template.validate_value(value)
+ hyperparameter[hyperparameter_name] = value
+ # use default value
+ else:
+ hyperparameter[hyperparameter_name] = hyperparameter_template.get_default_value()
+ return hyperparameter
+
+
+def generate_hyperparameters(
+ optuna_suggest: optuna.Trial,
+ task_options: TaskOptions
+) -> Tuple[Dict, Dict]:
+ """
+ Generate hyperparameters for model and task by optuna in auto-tuning training.
+ """
+ # TODO : support user to define hyperparameters in auto_tuning training
+ structure_hyperparameter_map = get_structure_hyperparameter_map(task_options.model_type)
+ model_hyperparameters = {}
+ for hyperparameter_name, hyperparameter_template in structure_hyperparameter_map.items():
+ model_hyperparameters[hyperparameter_name] = hyperparameter_template.suggest_parameter(optuna_suggest)
+
+ task_hyperparameters = {}
+ for hyperparameter_name, hyperparameter_template in training_hyperparameter_map.items():
+ task_hyperparameters[hyperparameter_name] = hyperparameter_template.suggest_parameter(optuna_suggest)
+ return model_hyperparameters, task_hyperparameters
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/metric.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/metric.py
new file mode 100644
index 00000000000..bcf63d1f68d
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/metric.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from abc import abstractmethod
+from typing import Dict, List
+
+import numpy as np
+
+all_metrics = ['RSE', 'CORR', 'MAE', 'MSE', 'RMSE', 'MAPE', 'MSPE']
+
+
+class Metric(object):
+ def __call__(self, pred, ground_truth):
+ return self.calculate(pred, ground_truth)
+
+ @abstractmethod
+ def calculate(self, pred, ground_truth):
+ pass
+
+
+class RSE(Metric):
+ def calculate(self, pred, ground_truth):
+ return np.sqrt(np.sum((ground_truth - pred) ** 2)) / np.sqrt(np.sum((ground_truth - ground_truth.mean()) ** 2))
+
+
+class CORR(Metric):
+ def calculate(self, pred, ground_truth):
+ u = ((ground_truth - ground_truth.mean(0)) * (pred - pred.mean(0))).sum(0)
+ d = np.sqrt(((ground_truth - ground_truth.mean(0)) ** 2 * (pred - pred.mean(0)) ** 2).sum(0))
+ return (u / d).mean(-1)
+
+
+class MAE(Metric):
+ def calculate(self, pred, ground_truth):
+ return np.mean(np.abs(pred - ground_truth))
+
+
+class MSE(Metric):
+ def calculate(self, pred, ground_truth):
+ return np.mean((pred - ground_truth) ** 2)
+
+
+class RMSE(Metric):
+ def calculate(self, pred, ground_truth):
+ mse = MSE()
+ return np.sqrt(mse(pred, ground_truth))
+
+
+class MAPE(Metric):
+ def calculate(self, pred, ground_truth):
+ return np.mean(np.abs((pred - ground_truth) / ground_truth))
+
+
+class MSPE(Metric):
+ def calculate(self, pred, ground_truth):
+ return np.mean(np.square((pred - ground_truth) / ground_truth))
+
+
+def build_metrics(metric_names: List) -> Dict[str, Metric]:
+ metrics_dict = {}
+ for metric_name in metric_names:
+ metrics_dict[metric_name] = eval(metric_name)()
+ return metrics_dict
+
+
+forecast_metric_names = ["MSE"]
diff --git a/iotdb-core/mlnode/requirements.txt b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/__init__.py
similarity index 88%
copy from iotdb-core/mlnode/requirements.txt
copy to iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/__init__.py
index c49c8a01896..2a1e720805f 100644
--- a/iotdb-core/mlnode/requirements.txt
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/__init__.py
@@ -15,13 +15,3 @@
# specific language governing permissions and limitations
# under the License.
#
-
-pandas>=1.3.5
-numpy>=1.21.4
-apache-iotdb
-poetry
-torch~=2.0.0
-pylru
-
-thrift~=0.13.0
-dynaconf~=3.1.12
\ No newline at end of file
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/dlinear.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/dlinear.py
new file mode 100644
index 00000000000..679c93d30c9
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/dlinear.py
@@ -0,0 +1,133 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import math
+from typing import Dict
+
+import torch
+import torch.nn as nn
+
+from iotdb.mlnode.constant import ModelInputName
+
+
+class MovingAverageBlock(nn.Module):
+ """ Moving average block to highlight the trend of time series """
+
+ def __init__(self, kernel_size, stride):
+ super(MovingAverageBlock, self).__init__()
+ self.kernel_size = kernel_size
+ self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+ def forward(self, x):
+ # padding on the both ends of time series
+ front = x[:, 0:1, :].repeat(1, self.kernel_size - 1 - math.floor((self.kernel_size - 1) // 2), 1)
+ end = x[:, -1:, :].repeat(1, math.floor((self.kernel_size - 1) // 2), 1)
+ x = torch.cat([front, x, end], dim=1)
+ x = self.avg(x.permute(0, 2, 1))
+ x = x.permute(0, 2, 1)
+ return x
+
+
+class SeriesDecompositionBlock(nn.Module):
+ """ Series decomposition block """
+
+ def __init__(self, kernel_size):
+ super(SeriesDecompositionBlock, self).__init__()
+ self.moving_avg = MovingAverageBlock(kernel_size, stride=1)
+
+ def forward(self, x):
+ moving_mean = self.moving_avg(x)
+ res = x - moving_mean
+ return res, moving_mean
+
+
+class DLinear(nn.Module):
+ """ Decomposition Linear Model """
+
+ def __init__(
+ self,
+ kernel_size: int,
+ input_len: int,
+ pred_len: int,
+ input_vars: int
+ ):
+ super(DLinear, self).__init__()
+ self.input_len = input_len
+ self.pred_len = pred_len
+ self.kernel_size = kernel_size
+ self.channels = input_vars
+
+ # decomposition Kernel Size
+ self.decomposition = SeriesDecompositionBlock(kernel_size)
+ self.linear_seasonal = nn.Linear(self.input_len, self.pred_len)
+ self.linear_trend = nn.Linear(self.input_len, self.pred_len)
+
+ def forward(self, input_dict: Dict):
+ # x: [Batch, Input length, Channel]
+ x = input_dict[ModelInputName.DATA_X.value]
+ seasonal_init, trend_init = self.decomposition(x)
+ seasonal_init, trend_init = seasonal_init.permute(0, 2, 1), trend_init.permute(0, 2, 1)
+
+ seasonal_output = self.linear_seasonal(seasonal_init)
+ trend_output = self.linear_trend(trend_init)
+
+ x = seasonal_output + trend_output
+ return x.permute(0, 2, 1) # to [Batch, Output length, Channel]
+
+
+class DLinearIndividual(nn.Module):
+ """ Decomposition Linear Model (individual) """
+
+ def __init__(
+ self,
+ kernel_size=25,
+ input_len=96,
+ pred_len=96,
+ input_vars=1
+ ):
+ super(DLinearIndividual, self).__init__()
+ self.input_len = input_len
+ self.pred_len = pred_len
+ self.kernel_size = kernel_size
+ self.channels = input_vars
+
+ self.decomposition = SeriesDecompositionBlock(kernel_size)
+ self.Linear_Seasonal = nn.ModuleList(
+ [nn.Linear(self.input_len, self.pred_len) for _ in range(self.channels)]
+ )
+ self.Linear_Trend = nn.ModuleList(
+ [nn.Linear(self.input_len, self.pred_len) for _ in range(self.channels)]
+ )
+
+ def forward(self, input_dict: Dict):
+ # x: [Batch, Input length, Channel]
+ x = input_dict[ModelInputName.DATA_X.value]
+ seasonal_init, trend_init = self.decomposition(x)
+ seasonal_init, trend_init = seasonal_init.permute(0, 2, 1), trend_init.permute(0, 2, 1)
+
+ seasonal_output = torch.zeros([seasonal_init.size(0), seasonal_init.size(1), self.pred_len],
+ dtype=seasonal_init.dtype).to(seasonal_init.device)
+ trend_output = torch.zeros([trend_init.size(0), trend_init.size(1), self.pred_len],
+ dtype=trend_init.dtype).to(trend_init.device)
+ for i, linear_season_layer in enumerate(self.Linear_Seasonal):
+ seasonal_output[:, i, :] = linear_season_layer(seasonal_init[:, i, :])
+ for i, linear_trend_layer in enumerate(self.Linear_Trend):
+ trend_output[:, i, :] = linear_trend_layer(trend_init[:, i, :])
+
+ x = seasonal_output + trend_output
+ return x.permute(0, 2, 1) # to [Batch, Output length, Channel]
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/nbeats.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/nbeats.py
new file mode 100644
index 00000000000..bf6d64528d4
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/models/forecast/nbeats.py
@@ -0,0 +1,131 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import Tuple, Dict
+
+import torch
+import torch.nn as nn
+
+from iotdb.mlnode.constant import ModelInputName
+
+
+class GenericBasis(nn.Module):
+ """ Generic basis function """
+
+ def __init__(self, backcast_size: int, forecast_size: int):
+ super().__init__()
+ self.backcast_size = backcast_size
+ self.forecast_size = forecast_size
+
+ def forward(self, theta: torch.Tensor):
+ return theta[:, :self.backcast_size], theta[:, -self.forecast_size:]
+
+
+block_dict = {
+ 'generic': GenericBasis,
+}
+
+
+class NBeatsBlock(nn.Module):
+ """ N-BEATS block which takes a basis function as an argument """
+
+ def __init__(self,
+ input_size,
+ theta_size: int,
+ basis_function: nn.Module,
+ layers: int,
+ layer_size: int):
+ """
+ N-BEATS block
+
+ Args:
+ input_size: input sample size
+ theta_size: number of parameters for the basis function
+ basis_function: basis function which takes the parameters and produces backcast and forecast
+ layers: number of layers
+ layer_size: layer size
+ """
+ super().__init__()
+ self.layers = nn.ModuleList([nn.Linear(in_features=input_size, out_features=layer_size)] + [
+ nn.Linear(in_features=layer_size, out_features=layer_size) for _ in range(layers - 1)])
+ self.basis_parameters = nn.Linear(in_features=layer_size, out_features=theta_size)
+ self.basis_function = basis_function
+
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ block_input = x
+ for layer in self.layers:
+ block_input = torch.relu(layer(block_input))
+ basis_parameters = self.basis_parameters(block_input)
+ return self.basis_function(basis_parameters)
+
+
+class NBeatsUnivariable(nn.Module):
+ """ N-Beats Model (uni-variable) """
+
+ def __init__(self, blocks: nn.ModuleList):
+ super().__init__()
+ self.blocks = blocks
+
+ def forward(self, x):
+ residuals = x
+ forecast = None
+ for _, block in enumerate(self.blocks):
+ backcast, block_forecast = block(residuals)
+ residuals = (residuals - backcast)
+ if forecast is None:
+ forecast = block_forecast
+ else:
+ forecast += block_forecast
+ return forecast
+
+
+class NBeats(nn.Module):
+ """ Neural Basis Expansion Analysis Time Series """
+
+ def __init__(
+ self,
+ block_type='generic',
+ d_model=128,
+ inner_layers=4,
+ outer_layers=4,
+ input_len=96,
+ pred_len=96,
+ input_vars=1
+ ):
+ super(NBeats, self).__init__()
+ self.enc_in = input_vars
+ self.block = block_dict[block_type]
+ self.model = NBeatsUnivariable(
+ torch.nn.ModuleList(
+ [NBeatsBlock(input_size=input_len,
+ theta_size=input_len + pred_len,
+ basis_function=self.block(backcast_size=input_len, forecast_size=pred_len),
+ layers=inner_layers,
+ layer_size=d_model)
+ for _ in range(outer_layers)]
+ )
+ )
+
+ def forward(self, input_dict: Dict):
+ # x: [Batch, Input length, Channel]
+ x = input_dict[ModelInputName.DATA_X.value]
+ res = []
+ for i in range(self.enc_in):
+ dec_out = self.model(x[:, :, i])
+ res.append(dec_out)
+ return torch.stack(res, dim=-1) # to [Batch, Output length, Channel]
diff --git a/iotdb-core/mlnode/iotdb/mlnode/algorithm/validator.py b/iotdb-core/mlnode/iotdb/mlnode/algorithm/validator.py
new file mode 100644
index 00000000000..1a671a5464e
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/algorithm/validator.py
@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from abc import abstractmethod
+
+
+class Validator(object):
+ @abstractmethod
+ def validate(self, value):
+ """
+ Checks whether the given value is valid.
+
+ Parameters:
+ - value: The value to validate
+
+ Returns:
+ - True if the value is valid, False otherwise.
+ """
+ raise NotImplementedError("Subclasses must implement the validate() method.")
+
+
+class NumberRangeValidator(Validator):
+ def __init__(self, min_value, max_value):
+ self.min_value = min_value
+ self.max_value = max_value
+
+ def validate(self, value):
+ if self.min_value <= value <= self.max_value:
+ return True
+ raise RuntimeError(
+ "Expect value between {0} and {1}, got {2} instead.".format(self.min_value, self.max_value, value))
diff --git a/iotdb-core/mlnode/iotdb/mlnode/client.py b/iotdb-core/mlnode/iotdb/mlnode/client.py
index 3157006e578..04013b9defa 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/client.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/client.py
@@ -19,7 +19,7 @@ import time
from typing import Dict, List
import pandas as pd
-from thrift.protocol import TBinaryProtocol, TCompactProtocol
+from thrift.protocol import TBinaryProtocol
from thrift.Thrift import TException
from thrift.transport import TSocket, TTransport
@@ -33,10 +33,9 @@ from iotdb.thrift.confignode import IConfigNodeRPCService
from iotdb.thrift.confignode.ttypes import (TUpdateModelInfoReq,
TUpdateModelStateReq)
from iotdb.thrift.datanode import IMLNodeInternalRPCService
-from iotdb.thrift.datanode.ttypes import (TFetchTimeseriesReq,
+from iotdb.thrift.datanode.ttypes import (TFetchMoreDataReq,
+ TFetchTimeseriesReq,
TRecordModelMetricsReq)
-from iotdb.thrift.mlnode import IMLNodeRPCService
-from iotdb.thrift.mlnode.ttypes import TCreateTrainingTaskReq, TDeleteModelReq
class ClientManager(object):
@@ -52,57 +51,6 @@ class ClientManager(object):
return ConfigNodeClient(config_leader=self.__config_node_endpoint)
-class MLNodeClient(object):
- def __init__(self, host, port):
- self.__host = host
- self.__port = port
-
- transport = TTransport.TBufferedTransport(
- TSocket.TSocket(self.__host, self.__port)
- )
- if not transport.isOpen():
- try:
- transport.open()
- except TTransport.TTransportException as e:
- logger.exception("TTransportException!", exc_info=e)
-
- protocol = TCompactProtocol.TCompactProtocol(transport)
- self.__client = IMLNodeRPCService.Client(protocol)
-
- def create_training_task(self,
- model_id: str,
- is_auto: bool,
- model_configs: Dict,
- query_expressions: List[str],
- query_filter: str = '') -> None:
- req = TCreateTrainingTaskReq(
- modelId=model_id,
- isAuto=is_auto,
- modelConfigs={k: str(v) for k, v in model_configs.items()},
- queryExpressions=[str(query) for query in query_expressions],
- queryFilter=query_filter,
- )
- try:
- status = self.__client.createTrainingTask(req)
- verify_success(status, "An error occurs when calling create_training_task()")
- except TTransport.TException as e:
- raise e
-
- def create_forecast_task(self) -> None:
- # TODO
- pass
-
- def delete_model(self,
- model_id: str,
- trial_id: str = None) -> None:
- req = TDeleteModelReq(modelId=model_id, trialId=trial_id)
- try:
- status = self.__client.deleteModel(req)
- verify_success(status, "An error occurs when calling delete_model()")
- except TTransport.TException as e:
- raise e
-
-
class DataNodeClient(object):
DEFAULT_FETCH_SIZE = 10000
DEFAULT_TIMEOUT = 60000
@@ -125,13 +73,11 @@ class DataNodeClient(object):
self.__client = IMLNodeInternalRPCService.Client(protocol)
def fetch_timeseries(self,
- query_expressions: List[str],
- query_filter: str = None,
+ query_body: str,
fetch_size: int = DEFAULT_FETCH_SIZE,
- timeout: int = DEFAULT_TIMEOUT) -> [int, bool, pd.DataFrame]:
+ timeout: int = DEFAULT_TIMEOUT) -> pd.DataFrame:
req = TFetchTimeseriesReq(
- queryExpressions=query_expressions,
- queryFilter=query_filter,
+ queryBody=query_body,
fetchSize=fetch_size,
timeout=timeout
)
@@ -140,27 +86,38 @@ class DataNodeClient(object):
verify_success(resp.status, "An error occurs when calling fetch_timeseries()")
if len(resp.tsDataset) == 0:
- raise RuntimeError(f'No data fetched with query filter: {query_filter}')
-
+ raise RuntimeError(f'No data fetched with sql: {query_body}')
data = serde.convert_to_df(resp.columnNameList,
resp.columnTypeList,
resp.columnNameIndexMap,
resp.tsDataset)
if data.empty:
raise RuntimeError(
- f'Fetched empty data with query expressions: {query_expressions} and query filter: {query_filter}')
- return resp.queryId, resp.hasMoreData, data
+ f'Fetched empty data with sql: {query_body}')
except Exception as e:
logger.warn(
- f'Fail to fetch data with query expressions: {query_expressions} and query filter: {query_filter}')
+ f'Fail to fetch data with sql: {query_body}')
raise e
-
- def fetch_window_batch(self,
- query_expressions: list,
- query_filter: str = None,
- fetch_size: int = DEFAULT_FETCH_SIZE,
- timeout: int = DEFAULT_TIMEOUT) -> [int, bool, List[pd.DataFrame]]:
- pass
+ query_id = resp.queryId
+ column_name_list = resp.columnNameList
+ column_type_list = resp.columnTypeList
+ column_name_index_map = resp.columnNameIndexMap
+ has_more_data = resp.hasMoreData
+ while has_more_data:
+ req = TFetchMoreDataReq(queryId=query_id, fetchSize=fetch_size)
+ try:
+ resp = self.__client.fetchMoreData(req)
+ verify_success(resp.status, "An error occurs when calling fetch_more_data()")
+ data = data.append(serde.convert_to_df(column_name_list,
+ column_type_list,
+ column_name_index_map,
+ resp.tsDataset))
+ has_more_data = resp.hasMoreData
+ except Exception as e:
+ logger.warn(
+ f'Fail to fetch more data with query id: {query_id}')
+ raise e
+ return data
def record_model_metrics(self,
model_id: str,
@@ -178,7 +135,7 @@ class DataNodeClient(object):
status = self.__client.recordModelMetrics(req)
verify_success(status, "An error occurs when calling record_model_metrics()")
except TTransport.TException as e:
- raise e
+ logger.exception(e.message)
class ConfigNodeClient(object):
@@ -271,9 +228,9 @@ class ConfigNodeClient(object):
req = TUpdateModelStateReq(
modelId=model_id,
state=training_state,
- bestTrailId=best_trail_id
+ bestTrialId=best_trail_id
)
- for i in range(0, self.__RETRY_NUM):
+ for _ in range(0, self.__RETRY_NUM):
try:
status = self.__client.updateModelState(req)
if not self.__update_config_node_leader(status):
@@ -295,11 +252,11 @@ class ConfigNodeClient(object):
model_info = {}
req = TUpdateModelInfoReq(
modelId=model_id,
- trailId=trial_id,
+ trialId=trial_id,
modelInfo={k: str(v) for k, v in model_info.items()},
)
- for i in range(0, self.__RETRY_NUM):
+ for _ in range(0, self.__RETRY_NUM):
try:
status = self.__client.updateModelInfo(req)
if not self.__update_config_node_leader(status):
diff --git a/iotdb-core/mlnode/iotdb/mlnode/config.py b/iotdb-core/mlnode/iotdb/mlnode/config.py
index 0ccfdc2cbb0..c54aa07ff25 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/config.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/config.py
@@ -40,6 +40,15 @@ class MLNodeConfig(object):
# Cache number of model storage to avoid repeated loading
self.__mn_model_storage_cache_size = 30
+ # Maximum number of training model tasks, otherwise the task is pending
+ self.__mn_task_pool_size = 10
+
+ # Maximum number of trials to be explored in a tuning task
+ self.__mn_tuning_trial_num = 20
+
+ # Concurrency of trials in a tuning task
+ self.__mn_tuning_trial_concurrency = 4
+
# Target ConfigNode to be connected by MLNode
self.__mn_target_config_node: TEndPoint = TEndPoint("127.0.0.1", 10710)
@@ -70,6 +79,24 @@ class MLNodeConfig(object):
def set_mn_model_storage_cache_size(self, mn_model_storage_cache_size: int) -> None:
self.__mn_model_storage_cache_size = mn_model_storage_cache_size
+ def get_mn_mn_task_pool_size(self) -> int:
+ return self.__mn_task_pool_size
+
+ def set_mn_task_pool_size(self, mn_task_pool_size: int) -> None:
+ self.__mn_task_pool_size = mn_task_pool_size
+
+ def get_mn_tuning_trial_num(self) -> int:
+ return self.__mn_tuning_trial_num
+
+ def set_mn_tuning_trial_num(self, mn_tuning_trial_num: int) -> None:
+ self.__mn_tuning_trial_num = mn_tuning_trial_num
+
+ def get_mn_tuning_trial_concurrency(self) -> int:
+ return self.__mn_tuning_trial_concurrency
+
+ def set_mn_tuning_trial_concurrency(self, mn_tuning_trial_concurrency: int) -> None:
+ self.__mn_tuning_trial_concurrency = mn_tuning_trial_concurrency
+
def get_mn_target_config_node(self) -> TEndPoint:
return self.__mn_target_config_node
@@ -114,6 +141,15 @@ class MLNodeDescriptor(object):
if file_configs.mn_model_storage_cache_size is not None:
self.__config.set_mn_model_storage_cache_size(file_configs.mn_model_storage_cache_size)
+ if file_configs.mn_task_pool_size is not None:
+ self.__config.set_mn_task_pool_size(file_configs.mn_task_pool_size)
+
+ if file_configs.mn_tuning_trial_num is not None:
+ self.__config.set_mn_tuning_trial_num(file_configs.mn_tuning_trial_num)
+
+ if file_configs.mn_tuning_trial_concurrency is not None:
+ self.__config.set_mn_tuning_trial_concurrency(file_configs.mn_tuning_trial_concurrency)
+
if file_configs.mn_target_config_node is not None:
self.__config.set_mn_target_config_node(file_configs.mn_target_config_node)
diff --git a/iotdb-core/mlnode/iotdb/mlnode/constant.py b/iotdb-core/mlnode/iotdb/mlnode/constant.py
index 68240af12a4..a3c25a6a7d4 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/constant.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/constant.py
@@ -16,13 +16,17 @@
# under the License.
#
from enum import Enum
+from typing import List
-MLNODE_CONF_DIRECTORY_NAME = "conf"
+MLNODE_CONF_DIRECTORY_NAME = "resources/conf"
MLNODE_CONF_FILE_NAME = "iotdb-mlnode.toml"
MLNODE_LOG_CONF_FILE_NAME = "logging_config.ini"
MLNODE_MODEL_STORAGE_DIRECTORY_NAME = "models"
+TRIAL_ID_PREFIX = "__trial_"
+DEFAULT_TRIAL_ID = TRIAL_ID_PREFIX + "0"
+
class TSStatusCode(Enum):
SUCCESS_STATUS = 200
@@ -31,3 +35,65 @@ class TSStatusCode(Enum):
def get_status_code(self) -> int:
return self.value
+
+
+class TaskType(Enum):
+ FORECAST = "forecast"
+
+
+class OptionsKey(Enum):
+ # common
+ TASK_TYPE = "task_type"
+ MODEL_TYPE = "model_type"
+ AUTO_TUNING = "auto_tuning"
+ INPUT_VARS = "input_vars"
+
+ # forecast
+ INPUT_LENGTH = "input_length"
+ PREDICT_LENGTH = "predict_length"
+ PREDICT_INDEX_LIST = "predict_index_list"
+ INPUT_TYPE_LIST = "input_type_list"
+
+ def name(self) -> str:
+ return self.value
+
+
+class HyperparameterName(Enum):
+ # Training hyperparameter
+ LEARNING_RATE = "learning_rate"
+ EPOCHS = "epochs"
+ BATCH_SIZE = "batch_size"
+ USE_GPU = "use_gpu"
+ NUM_WORKERS = "num_workers"
+
+ # Structure hyperparameter
+ KERNEL_SIZE = "kernel_size"
+ INPUT_VARS = "input_vars"
+ BLOCK_TYPE = "block_type"
+ D_MODEL = "d_model"
+ INNER_LAYERS = "inner_layer"
+ OUTER_LAYERS = "outer_layer"
+
+ def name(self):
+ return self.value
+
+
+class ForecastModelType(Enum):
+ DLINEAR = "dlinear"
+ DLINEAR_INDIVIDUAL = "dlinear_individual"
+ NBEATS = "nbeats"
+
+ @classmethod
+ def values(cls) -> List[str]:
+ values = []
+ for item in list(cls):
+ values.append(item.value)
+ return values
+
+
+class ModelInputName(Enum):
+ DATA_X = "data_x"
+ TIME_STAMP_X = "time_stamp_x"
+ TIME_STAMP_Y = "time_stamp_y"
+ DEC_INP = "dec_inp"
+
diff --git a/iotdb-core/mlnode/iotdb/mlnode/dataset/dataset.py b/iotdb-core/mlnode/iotdb/mlnode/dataset/dataset.py
new file mode 100644
index 00000000000..ab9a1c1c5b3
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/dataset/dataset.py
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from typing import Tuple
+
+import numpy as np
+from torch.utils.data import Dataset
+
+from iotdb.mlnode.dataset.source import DataSource
+from iotdb.mlnode.dataset.utils.time_features import time_features
+
+
+class TsForecastDataset(Dataset):
+ """
+ Build Windowed dataset (with each element as multivariable time series
+ with a sliding window and corresponding timestamps embedding),
+ the sliding step is one unit in give data source
+
+ Args:
+ data_source: the whole multivariate time series for a while
+ time_embed: embedding frequency, see `utils/time_features.py` for more detail
+ input_len: input window size (unit) [1, 2, ... I]
+ pred_len: output window size (unit) right after the input window [I+1, I+2, ... I+P]
+
+ Returns:
+ Random accessible dataset
+ """
+
+ def __init__(self,
+ data_source: DataSource,
+ input_len: int,
+ pred_len: int,
+ time_embed: str = 'h'):
+ self.input_len = input_len
+ self.pred_len = pred_len
+ self.time_embed = time_embed
+ self.data = data_source.get_data()
+ self.data_stamp = time_features(data_source.get_timestamp(), time_embed=self.time_embed).transpose(1, 0)
+ self.n_vars = self.data.shape[-1]
+ if input_len > self.data.shape[0]:
+ raise RuntimeError('input_len should not be larger than the number of time series points')
+ if pred_len > self.data.shape[0]:
+ raise RuntimeError('pred_len should not be larger than the number of time series points')
+
+ def __getitem__(self, index) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+ s_begin = index
+ s_end = s_begin + self.input_len
+ r_begin = s_end
+ r_end = s_end + self.pred_len
+ seq_x = self.data[s_begin:s_end]
+ seq_y = self.data[r_begin:r_end]
+ seq_x_t = self.data_stamp[s_begin:s_end]
+ seq_y_t = self.data_stamp[r_begin:r_end]
+ return seq_x, seq_y, seq_x_t, seq_y_t
+
+ def __len__(self) -> int:
+ return len(self.data) - self.input_len - self.pred_len + 1
+
+ def get_variable_num(self) -> int:
+ return self.n_vars # number of series in data_source
diff --git a/iotdb-core/mlnode/iotdb/mlnode/dataset/factory.py b/iotdb-core/mlnode/iotdb/mlnode/dataset/factory.py
new file mode 100644
index 00000000000..de666c52206
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/dataset/factory.py
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from typing import cast
+
+from torch.utils.data import Dataset
+
+from iotdb.mlnode.constant import TaskType
+from iotdb.mlnode.dataset.dataset import TsForecastDataset
+from iotdb.mlnode.dataset.source import IoTDBDataSource
+from iotdb.mlnode.exception import UnsupportedError
+from iotdb.mlnode.parser import ForecastTaskOptions, TaskOptions
+
+
+def create_dataset(query_body: str, task_options: TaskOptions) -> Dataset:
+ """
+ Create a dataset for training according to the task type.
+ Currently, only ForcastTask is supported, so only create_forecast_dataset is called.
+ """
+ task_type = task_options.get_task_type()
+ if task_type == TaskType.FORECAST:
+ task_options = cast(ForecastTaskOptions, task_options)
+ return create_forecast_dataset(query_body, task_options)
+ else:
+ raise UnsupportedError(f"task type {task_type}")
+
+
+def create_forecast_dataset(query_body: str, task_options: ForecastTaskOptions) -> TsForecastDataset:
+ """
+ Create a dataset for training for forecasting.
+ In forecasting, dataset is created from a IoTDBDataSource, and then wrapped by a TsDataset,
+ where the input length and predict length are required.
+ """
+ datasource = IoTDBDataSource(query_body)
+ dataset = TsForecastDataset(datasource, task_options.input_length, task_options.predict_length)
+ return dataset
diff --git a/iotdb-core/mlnode/iotdb/mlnode/dataset/source.py b/iotdb-core/mlnode/iotdb/mlnode/dataset/source.py
new file mode 100644
index 00000000000..f69e7854374
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/dataset/source.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import numpy as np
+import pandas as pd
+
+from iotdb.mlnode.client import client_manager
+
+
+class DataSource(object):
+ """
+ Pre-fetched in multi-variate time series in memory
+
+ Methods:
+ get_data: returns time series value (Numpy.2DArray)
+ get_timestamp: returns aligned timestamp value
+ """
+
+ def __init__(self):
+ self.data = None
+ self.timestamp = None
+ self._read_data()
+
+ def _read_data(self):
+ raise NotImplementedError
+
+ def get_data(self) -> np.ndarray:
+ return self.data
+
+ def get_timestamp(self) -> np.ndarray:
+ return self.timestamp
+
+
+class FileDataSource(DataSource):
+ def __init__(self, filename: str = None):
+ self.filename = filename
+ super(FileDataSource, self).__init__()
+
+ def _read_data(self) -> None:
+ try:
+ raw_data = pd.read_csv(self.filename)
+ except Exception:
+ raise RuntimeError(f'Fail to load data with filename: {self.filename}')
+ cols_data = raw_data.columns[1:]
+ self.data = raw_data[cols_data].values
+ self.timestamp = pd.to_datetime(raw_data[raw_data.columns[0]].values)
+
+
+class IoTDBDataSource(DataSource):
+ def __init__(self, query_body: str = None):
+ self.query_body = query_body
+ super(IoTDBDataSource, self).__init__()
+
+ def _read_data(self) -> None:
+ try:
+ data_client = client_manager.borrow_data_node_client()
+ except Exception:
+ raise RuntimeError('Fail to establish connection with DataNode')
+
+ raw_data = data_client.fetch_timeseries(self.query_body)
+
+ cols_data = raw_data.columns[1:]
+ self.data = raw_data[cols_data].values
+ self.timestamp = pd.to_datetime(raw_data[raw_data.columns[0]].values, unit='ms', utc=True) \
+ .tz_convert('Asia/Shanghai')
diff --git a/iotdb-core/mlnode/requirements.txt b/iotdb-core/mlnode/iotdb/mlnode/dataset/utils/__init__.py
similarity index 88%
copy from iotdb-core/mlnode/requirements.txt
copy to iotdb-core/mlnode/iotdb/mlnode/dataset/utils/__init__.py
index c49c8a01896..2a1e720805f 100644
--- a/iotdb-core/mlnode/requirements.txt
+++ b/iotdb-core/mlnode/iotdb/mlnode/dataset/utils/__init__.py
@@ -15,13 +15,3 @@
# specific language governing permissions and limitations
# under the License.
#
-
-pandas>=1.3.5
-numpy>=1.21.4
-apache-iotdb
-poetry
-torch~=2.0.0
-pylru
-
-thrift~=0.13.0
-dynaconf~=3.1.12
\ No newline at end of file
diff --git a/iotdb-core/mlnode/iotdb/mlnode/dataset/utils/time_features.py b/iotdb-core/mlnode/iotdb/mlnode/dataset/utils/time_features.py
new file mode 100644
index 00000000000..17e5a18267a
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/dataset/utils/time_features.py
@@ -0,0 +1,149 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from typing import List
+
+import numpy as np
+import pandas as pd
+from pandas.tseries import offsets
+from pandas.tseries.frequencies import to_offset
+
+
+class TimeFeature:
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ raise NotImplementedError
+
+ def __repr__(self):
+ return self.__class__.__name__ + '()'
+
+
+class SecondOfMinute(TimeFeature):
+ """Minute of hour encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return index.second / 59.0 - 0.5
+
+
+class MinuteOfHour(TimeFeature):
+ """Minute of hour encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return index.minute / 59.0 - 0.5
+
+
+class HourOfDay(TimeFeature):
+ """Hour of day encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return index.hour / 23.0 - 0.5
+
+
+class DayOfWeek(TimeFeature):
+ """Hour of day encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return index.dayofweek / 6.0 - 0.5
+
+
+class DayOfMonth(TimeFeature):
+ """Day of month encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return (index.day - 1) / 30.0 - 0.5
+
+
+class DayOfYear(TimeFeature):
+ """Day of year encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return (index.dayofyear - 1) / 365.0 - 0.5
+
+
+class MonthOfYear(TimeFeature):
+ """Month of year encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return (index.month - 1) / 11.0 - 0.5
+
+
+class WeekOfYear(TimeFeature):
+ """Week of year encoded as value between [-0.5, 0.5]"""
+
+ def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
+ return (index.isocalendar().week - 1) / 52.0 - 0.5
+
+
+def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:
+ """
+ Embedding timestamp by given frequency string
+ Args:
+ freq_str: frequency string of the form [multiple][granularity] such as '12H', '5min', '1D' etc.
+ Returns:
+ a list of time features that will be appropriate for the given frequency string.
+ """
+
+ features_by_offsets = {
+ offsets.YearEnd: [],
+ offsets.QuarterEnd: [MonthOfYear],
+ offsets.MonthEnd: [MonthOfYear],
+ offsets.Week: [DayOfMonth, WeekOfYear],
+ offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],
+ offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],
+ offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],
+ offsets.Minute: [
+ MinuteOfHour,
+ HourOfDay,
+ DayOfWeek,
+ DayOfMonth,
+ DayOfYear,
+ ],
+ offsets.Second: [
+ SecondOfMinute,
+ MinuteOfHour,
+ HourOfDay,
+ DayOfWeek,
+ DayOfMonth,
+ DayOfYear,
+ ],
+ }
+
+ try:
+ offset = to_offset(freq_str)
+
+ for offset_type, feature_classes in features_by_offsets.items():
+ if isinstance(offset, offset_type):
+ return [cls() for cls in feature_classes]
+ except ValueError:
+ supported_freq_msg = f'''
+ Unsupported time embedding frequency ({freq_str})
+ The following frequencies are supported (case-insensitive):
+ Y - yearly
+ alias: A
+ M - monthly
+ W - weekly
+ D - daily
+ B - business days
+ H - hourly
+ T - minutely
+ alias: min
+ S - secondly
+ '''
+ raise RuntimeError(supported_freq_msg)
+
+
+def time_features(dates, time_embed='h'):
+ return np.vstack([feat(dates) for feat in time_features_from_frequency_str(time_embed)])
diff --git a/iotdb-core/mlnode/iotdb/mlnode/exception.py b/iotdb-core/mlnode/iotdb/mlnode/exception.py
index 6307909a9ac..36beab83307 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/exception.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/exception.py
@@ -18,7 +18,12 @@
class _BaseError(Exception):
"""Base class for exceptions in this module."""
- pass
+
+ def __init__(self):
+ self.message = None
+
+ def __str__(self) -> str:
+ return self.message
class BadNodeUrlError(_BaseError):
@@ -28,4 +33,34 @@ class BadNodeUrlError(_BaseError):
class ModelNotExistError(_BaseError):
def __init__(self, file_path: str):
- self.message = "Model path: ({}) not exists".format(file_path)
+ self.message = "Model path is not exists: {} ".format(file_path)
+
+
+class BadConfigValueError(_BaseError):
+ def __init__(self, config_name: str, config_value, hint: str = ''):
+ self.message = "Bad value [{0}] for config {1}. {2}".format(config_value, config_name, hint)
+
+
+class MissingConfigError(_BaseError):
+ def __init__(self, config_name: str):
+ self.message = "Missing config: {}".format(config_name)
+
+
+class MissingOptionError(_BaseError):
+ def __init__(self, config_name: str):
+ self.message = "Missing task option: {}".format(config_name)
+
+
+class RedundantOptionError(_BaseError):
+ def __init__(self, option_name: str):
+ self.message = "Redundant task option: {}".format(option_name)
+
+
+class WrongTypeConfigError(_BaseError):
+ def __init__(self, config_name: str, expected_type: str):
+ self.message = "Wrong type for config: {0}, expected: {1}".format(config_name, expected_type)
+
+
+class UnsupportedError(_BaseError):
+ def __init__(self, msg: str):
+ self.message = "{0} is not supported in current version".format(msg)
diff --git a/iotdb-core/mlnode/iotdb/mlnode/handler.py b/iotdb-core/mlnode/iotdb/mlnode/handler.py
index 78021420d43..0790f58b606 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/handler.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/handler.py
@@ -15,7 +15,18 @@
# specific language governing permissions and limitations
# under the License.
#
-from iotdb.mlnode.constant import TSStatusCode
+
+from typing import cast
+
+from iotdb.mlnode.config import descriptor
+from iotdb.mlnode.constant import TaskType, TSStatusCode
+from iotdb.mlnode.dataset.dataset import TsForecastDataset
+from iotdb.mlnode.dataset.factory import create_dataset
+from iotdb.mlnode.log import logger
+from iotdb.mlnode.parser import (ForecastTaskOptions, parse_forecast_request,
+ parse_task_options)
+from iotdb.mlnode.process.manager import TaskManager
+from iotdb.mlnode.serde import convert_to_binary
from iotdb.mlnode.storage import model_storage
from iotdb.mlnode.util import get_status
from iotdb.thrift.mlnode import IMLNodeRPCService
@@ -26,19 +37,59 @@ from iotdb.thrift.mlnode.ttypes import (TCreateTrainingTaskReq,
class MLNodeRPCServiceHandler(IMLNodeRPCService.Iface):
def __init__(self):
- pass
+ self.__task_manager = TaskManager(pool_size=descriptor.get_config().get_mn_mn_task_pool_size())
def deleteModel(self, req: TDeleteModelReq):
try:
model_storage.delete_model(req.modelId)
return get_status(TSStatusCode.SUCCESS_STATUS)
except Exception as e:
+ logger.warn(e)
return get_status(TSStatusCode.MLNODE_INTERNAL_ERROR, str(e))
def createTrainingTask(self, req: TCreateTrainingTaskReq):
- return get_status(TSStatusCode.SUCCESS_STATUS)
+ task = None
+ try:
+ # parse options
+ task_options = parse_task_options(req.options)
+ dataset = create_dataset(req.datasetFetchSQL, task_options)
+
+ # create task according to task type
+ # currently, IoTDB-ML supports forecasting training task only
+ if task_options.get_task_type() == TaskType.FORECAST:
+ task_options = cast(ForecastTaskOptions, task_options)
+ dataset = cast(TsForecastDataset, dataset)
+ task = self.__task_manager.create_forecast_training_task(
+ model_id=req.modelId,
+ task_options=task_options,
+ hyperparameters=req.hyperparameters,
+ dataset=dataset
+ )
+ else:
+ raise NotImplementedError
+
+ return get_status(TSStatusCode.SUCCESS_STATUS)
+ except Exception as e:
+ logger.warn(e)
+ return get_status(TSStatusCode.MLNODE_INTERNAL_ERROR, str(e))
+ finally:
+ if task is not None:
+ # submit task to process pool
+ self.__task_manager.submit_training_task(task)
def forecast(self, req: TForecastReq):
- status = get_status(TSStatusCode.SUCCESS_STATUS)
- forecast_result = b'forecast result'
- return TForecastResp(status, forecast_result)
+ model_path, data, pred_length = parse_forecast_request(req)
+ try:
+ task = self.__task_manager.create_forecast_task(
+ pred_length,
+ data,
+ model_path
+ )
+ # submit task stage & check resource and decide pending/start
+ res = self.__task_manager.submit_forecast_task(task)
+ forecast_result = convert_to_binary(res)
+ resp = TForecastResp(get_status(TSStatusCode.SUCCESS_STATUS), forecast_result)
+ return resp
+ except Exception as e:
+ logger.warn(e)
+ return get_status(TSStatusCode.MLNODE_INTERNAL_ERROR, str(e))
diff --git a/iotdb-core/mlnode/iotdb/mlnode/parser.py b/iotdb-core/mlnode/iotdb/mlnode/parser.py
new file mode 100644
index 00000000000..8299ac6f658
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/parser.py
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+import argparse
+from abc import abstractmethod
+from typing import Dict
+
+from iotdb.mlnode.constant import OptionsKey, TaskType, ForecastModelType
+from iotdb.mlnode.exception import (MissingOptionError, RedundantOptionError,
+ UnsupportedError)
+from iotdb.mlnode.serde import convert_to_df
+from iotdb.thrift.mlnode.ttypes import TForecastReq
+
+
+class TaskOptions(object):
+ def __init__(self, options: Dict):
+ self._raw_options = options
+
+ if OptionsKey.MODEL_TYPE.name() not in self._raw_options:
+ raise MissingOptionError(OptionsKey.MODEL_TYPE.name())
+ model_name = self._raw_options.pop(OptionsKey.MODEL_TYPE.name())
+ self.model_type = getattr(ForecastModelType, model_name.upper(), None)
+ if not self.model_type:
+ raise UnsupportedError(f"model_type {model_name}")
+
+ # training with auto-tuning as default
+ self.auto_tuning = str2bool(self._raw_options.pop(OptionsKey.AUTO_TUNING.name(), "false"))
+
+ @abstractmethod
+ def get_task_type(self) -> TaskType:
+ raise NotImplementedError("Subclasses must implement the validate() method.")
+
+ def _check_redundant_options(self) -> None:
+ if len(self._raw_options):
+ raise RedundantOptionError(str(self._raw_options))
+
+
+class ForecastTaskOptions(TaskOptions):
+ def __init__(self, options: Dict):
+ super().__init__(options)
+ self.input_length = self._raw_options.pop(OptionsKey.INPUT_LENGTH.name(), 96)
+ self.predict_length = self._raw_options.pop(OptionsKey.PREDICT_LENGTH.name(), 96)
+ self.predict_index_list = self._raw_options.pop(OptionsKey.PREDICT_INDEX_LIST.name(), None)
+ self.input_type_list = self._raw_options.pop(OptionsKey.INPUT_TYPE_LIST.name(), None)
+ super()._check_redundant_options()
+
+ def get_task_type(self) -> TaskType:
+ return TaskType.FORECAST
+
+
+def parse_task_type(options: Dict) -> TaskType:
+ if OptionsKey.TASK_TYPE.name() not in options:
+ raise MissingOptionError(OptionsKey.TASK_TYPE.name())
+ task_name = options.pop(OptionsKey.TASK_TYPE.name())
+ task_type = getattr(TaskType, task_name.upper(), None)
+ if not task_type:
+ raise UnsupportedError(f"task_type {task_name}")
+ return task_type
+
+
+def parse_task_options(options) -> TaskOptions:
+ task_type = parse_task_type(options)
+ if task_type == TaskType.FORECAST:
+ return ForecastTaskOptions(options)
+ else:
+ raise UnsupportedError(f"task type {task_type}")
+
+
+def parse_forecast_request(req: TForecastReq):
+ column_name_list = req.inputColumnNameList
+ column_type_list = req.inputTypeList
+ ts_dataset = req.inputData
+ data = convert_to_df(column_name_list, column_type_list, None, [ts_dataset])
+ time_stamp, data = data[data.columns[0:1]], data[data.columns[1:]]
+ full_data = (data, time_stamp)
+ return req.modelPath, full_data, req.predictLength
+
+
+def str2bool(value):
+ if value.lower() in ('yes', 'true', 't', 'y', '1'):
+ return True
+ elif value.lower() in ('no', 'false', 'f', 'n', '0'):
+ return False
+ else:
+ raise argparse.ArgumentTypeError('Boolean value expected.')
diff --git a/iotdb-core/mlnode/requirements.txt b/iotdb-core/mlnode/iotdb/mlnode/process/__init__.py
similarity index 88%
copy from iotdb-core/mlnode/requirements.txt
copy to iotdb-core/mlnode/iotdb/mlnode/process/__init__.py
index c49c8a01896..2a1e720805f 100644
--- a/iotdb-core/mlnode/requirements.txt
+++ b/iotdb-core/mlnode/iotdb/mlnode/process/__init__.py
@@ -15,13 +15,3 @@
# specific language governing permissions and limitations
# under the License.
#
-
-pandas>=1.3.5
-numpy>=1.21.4
-apache-iotdb
-poetry
-torch~=2.0.0
-pylru
-
-thrift~=0.13.0
-dynaconf~=3.1.12
\ No newline at end of file
diff --git a/iotdb-core/mlnode/iotdb/mlnode/process/manager.py b/iotdb-core/mlnode/iotdb/mlnode/process/manager.py
new file mode 100644
index 00000000000..83c929f4993
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/process/manager.py
@@ -0,0 +1,124 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import multiprocessing as mp
+import signal
+import sys
+from subprocess import call
+from typing import Dict
+
+import pandas as pd
+import psutil
+
+from iotdb.mlnode.dataset.dataset import TsForecastDataset
+from iotdb.mlnode.log import logger
+from iotdb.mlnode.parser import ForecastTaskOptions
+from iotdb.mlnode.process.task import (ForecastAutoTuningTrainingTask,
+ ForecastFixedParamTrainingTask,
+ ForecastingInferenceTask,
+ _BasicTrainingTask)
+
+
+class TaskManager(object):
+ def __init__(self, pool_size: int):
+ """
+ Args:
+ pool_size: specify the maximum process number of the process pool
+
+ __shared_resource_manager: a manager that manage resources shared between processes
+ __pid_info: a map shared between processes, can be used to find the pid with model_id and trial_id
+ __training_process_pool: a multiprocessing process pool
+ """
+ self.__shared_resource_manager = mp.Manager()
+ self.__pid_info = self.__shared_resource_manager.dict()
+ self.__training_process_pool = mp.Pool(pool_size)
+ self.__inference_process_pool = mp.Pool(pool_size)
+
+ def create_forecast_training_task(self,
+ model_id: str,
+ task_options: ForecastTaskOptions,
+ hyperparameters: Dict[str, str],
+ dataset: TsForecastDataset):
+ """
+ Create a training task for forecasting, which contains the training process
+
+ Args:
+ model_id: the unique id of the model
+ task_options: the options of the task, contains task_type, model_type, etc.
+ dataset: a torch dataset to be used for training
+ hyperparameters: a dict of hyperparameters, which is consisted of model_hyperparameters
+ and task_hyperparameters
+
+ Returns:
+ task: a training task for forecasting, which can be submitted to self.__training_process_pool
+ """
+ if task_options.auto_tuning:
+ return ForecastAutoTuningTrainingTask(
+ model_id=model_id,
+ task_options=task_options,
+ hyperparameters=hyperparameters,
+ dataset=dataset,
+ pid_info=self.__pid_info,
+ )
+ else:
+ return ForecastFixedParamTrainingTask(
+ model_id=model_id,
+ task_options=task_options,
+ hyperparameters=hyperparameters,
+ dataset=dataset,
+ pid_info=self.__pid_info,
+ )
+
+ def submit_training_task(self, task: _BasicTrainingTask) -> None:
+ self.__training_process_pool.apply_async(task, args=())
+ logger.info(f'Task: ({task.model_id}) - Training process submitted successfully')
+
+ def create_forecast_task(self,
+ predict_length,
+ data,
+ model_path) -> ForecastingInferenceTask:
+ task = ForecastingInferenceTask(
+ predict_length,
+ self.__pid_info,
+ data,
+ model_path
+ )
+ return task
+
+ def submit_forecast_task(self, task: ForecastingInferenceTask) -> pd.DataFrame:
+ read_pipe, send_pipe = mp.Pipe()
+ if task is not None:
+ self.__inference_process_pool.apply_async(task, args=(send_pipe,))
+ logger.info('Forecasting process submitted successfully')
+ return read_pipe.recv()
+
+ def kill_task(self, model_id):
+ """
+ Kill the process by pid, will check whether the pid is training or inference process
+ """
+ pid = self.__pid_info[model_id]
+ if sys.platform == 'win32':
+ try:
+ process = psutil.Process(pid=pid)
+ process.send_signal(signal.CTRL_BREAK_EVENT)
+ except psutil.NoSuchProcess:
+ print(f'Tried to kill process (pid = {pid}), '
+ f'but the process does not exist.')
+ else:
+ cmds = ['kill', str(pid)]
+ call(cmds)
diff --git a/iotdb-core/mlnode/iotdb/mlnode/process/task.py b/iotdb-core/mlnode/iotdb/mlnode/process/task.py
new file mode 100644
index 00000000000..dd71a0b3a27
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/process/task.py
@@ -0,0 +1,265 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import os
+from abc import abstractmethod
+from multiprocessing.connection import Connection
+from typing import Dict, Tuple
+
+import numpy as np
+import optuna
+import pandas as pd
+import torch
+
+from iotdb.mlnode.algorithm.hyperparameter import (generate_hyperparameters,
+ parse_fixed_hyperparameters)
+from iotdb.mlnode.client import client_manager
+from iotdb.mlnode.config import descriptor
+from iotdb.mlnode.constant import DEFAULT_TRIAL_ID, TRIAL_ID_PREFIX, OptionsKey, ModelInputName
+from iotdb.mlnode.dataset.dataset import TsForecastDataset
+from iotdb.mlnode.log import logger
+from iotdb.mlnode.parser import ForecastTaskOptions
+from iotdb.mlnode.process.trial import ForecastingTrainingTrial
+from iotdb.mlnode.storage import model_storage
+from iotdb.mlnode.util import pack_input_dict
+from iotdb.thrift.common.ttypes import TrainingState
+
+
+class _BasicTask(object):
+ def __init__(
+ self,
+ pid_info: Dict
+ ):
+ self.pid_info = pid_info
+
+ @abstractmethod
+ def __call__(self):
+ raise NotImplementedError
+
+
+# Training Task
+
+class _BasicTrainingTask(_BasicTask):
+ def __init__(
+ self,
+ model_id: str,
+ hyperparameters: Dict[str, str],
+ pid_info: Dict
+ ):
+ super().__init__(pid_info)
+ self.model_id = model_id
+ self.hyperparameters = hyperparameters
+ self.configNode_client = client_manager.borrow_config_node_client()
+
+ @abstractmethod
+ def __call__(self):
+ raise NotImplementedError
+
+
+class ForecastFixedParamTrainingTask(_BasicTrainingTask):
+ def __init__(
+ self,
+ model_id: str,
+ task_options: ForecastTaskOptions,
+ hyperparameters: Dict[str, str],
+ dataset: TsForecastDataset,
+ pid_info: Dict
+ ):
+ super().__init__(model_id, hyperparameters, pid_info)
+ model_hyperparameters, task_hyperparameters = parse_fixed_hyperparameters(task_options, hyperparameters)
+ self.dataset = dataset
+ self.trial = ForecastingTrainingTrial(
+ trial_id=DEFAULT_TRIAL_ID,
+ model_id=model_id,
+ task_options=task_options,
+ model_hyperparameters=model_hyperparameters,
+ task_hyperparameters=task_hyperparameters,
+ dataset=dataset)
+
+ def __call__(self):
+ try:
+ self.configNode_client.update_model_state(self.model_id, TrainingState.RUNNING)
+ self.pid_info[self.model_id] = os.getpid()
+ self.trial.start()
+ self.configNode_client.update_model_state(self.model_id, TrainingState.FINISHED, self.trial.trial_id)
+ except Exception as e:
+ logger.warn(e)
+ self.configNode_client.update_model_state(self.model_id, TrainingState.FAILED, self.trial.trial_id)
+ raise e
+
+
+class ForestingTrainingObjective:
+ """
+ A class which serve as a function, should accept trial as args
+ and return the optimization objective.
+ Optuna will try to minimize the objective.
+ Currently, user cannot define the range of hyperparameters.
+ """
+
+ def __init__(
+ self,
+ model_id: str,
+ task_options: ForecastTaskOptions,
+ hyperparameters: Dict[str, str],
+ dataset: TsForecastDataset,
+ pid_info: Dict
+ ):
+ self.task_options = task_options
+ self.hyperparameters = hyperparameters
+ self.dataset = dataset
+ self.pid_info = pid_info
+ self.model_id = model_id
+
+ def __call__(self, optuna_suggest: optuna.Trial):
+ model_hyperparameters, task_hyperparameters = generate_hyperparameters(optuna_suggest,
+ self.task_options)
+ trial = ForecastingTrainingTrial(
+ trial_id=TRIAL_ID_PREFIX + str(optuna_suggest._trial_id),
+ model_id=self.model_id,
+ task_options=self.task_options,
+ model_hyperparameters=model_hyperparameters,
+ task_hyperparameters=task_hyperparameters,
+ dataset=self.dataset)
+ loss = trial.start()
+ return loss
+
+
+class ForecastAutoTuningTrainingTask(_BasicTrainingTask):
+ def __init__(
+ self,
+ model_id: str,
+ task_options: ForecastTaskOptions,
+ hyperparameters: Dict[str, str],
+ dataset: TsForecastDataset,
+ pid_info: Dict
+ ):
+ super().__init__(model_id, hyperparameters, pid_info)
+ self.study = optuna.create_study(direction='minimize')
+ self.task_options = task_options
+ self.dataset = dataset
+
+ def __call__(self):
+ self.pid_info[self.model_id] = os.getpid()
+ try:
+ self.configNode_client.update_model_state(self.model_id, TrainingState.RUNNING)
+ self.study.optimize(ForestingTrainingObjective(
+ model_id=self.model_id,
+ task_options=self.task_options,
+ hyperparameters=self.hyperparameters,
+ dataset=self.dataset,
+ pid_info=self.pid_info),
+ n_trials=descriptor.get_config().get_mn_tuning_trial_num(),
+ n_jobs=descriptor.get_config().get_mn_tuning_trial_concurrency())
+ best_trial_id = TRIAL_ID_PREFIX + str(self.study.best_trial._trial_id)
+ self.configNode_client.update_model_state(self.model_id, TrainingState.FINISHED, best_trial_id)
+ except Exception as e:
+ logger.warn(e)
+ self.configNode_client.update_model_state(self.model_id, TrainingState.FAILED)
+ raise e
+
+
+# Inference Task
+
+class _BasicInferenceTask(_BasicTask):
+ def __init__(
+ self,
+ pid_info: Dict,
+ data: Tuple,
+ model_path: str
+ ):
+ super().__init__(pid_info)
+ self.data = data
+ self.model_path = model_path
+
+ @abstractmethod
+ def __call__(self, pipe: Connection = None):
+ raise NotImplementedError
+
+ @abstractmethod
+ def data_align(self, *args):
+ raise NotImplementedError
+
+
+class ForecastingInferenceTask(_BasicInferenceTask):
+ def __init__(
+ self,
+ predict_length: int,
+ pid_info: Dict,
+ data: Tuple,
+ model_path: str
+ ):
+ super().__init__(pid_info, data, model_path)
+ self.pred_len = predict_length
+
+ def __call__(self, pipe: Connection = None):
+ self.model, self.model_configs = model_storage.load_model(self.model_path)
+ self.model_pred_len = self.model_configs[OptionsKey.PREDICT_LENGTH.name()]
+ self.model_input_len = self.model_configs[OptionsKey.INPUT_LENGTH.name()]
+
+ data, time_stamp = self.data
+ _, c = data.shape
+
+ time_stamp = pd.to_datetime(time_stamp.values[:, 0], unit='ms', utc=True) \
+ .tz_convert('Asia/Shanghai') # for iotdb
+ data, time_stamp = self.data_align(data, time_stamp)
+ full_data, full_data_stamp = data, time_stamp
+ current_pred_len = 0
+ while current_pred_len < self.pred_len:
+ current_data = full_data[:, -self.model_input_len:, :]
+ current_data = torch.Tensor(current_data)
+
+ input_data = pack_input_dict(batch_x=current_data)
+ output_data = self.model(input_data).detach().numpy()
+ full_data = np.concatenate([full_data, output_data], axis=1)
+ current_pred_len += self.model_pred_len
+ full_data_stamp = self.generate_future_mark(full_data_stamp, self.pred_len)
+ ret_data = pd.concat(
+ [pd.DataFrame(full_data_stamp.astype(np.int64)),
+ pd.DataFrame(full_data[0, -self.pred_len:, :]).astype(np.double)], axis=1)
+ ret_data.columns = list(np.arange(0, c + 1))
+ pipe.send(ret_data)
+
+ def data_align(self, data: pd.DataFrame, data_stamp) -> Tuple[np.ndarray, pd.DataFrame]:
+ """
+ data: L x C, DataFrame, suppose no batch dim
+ time_stamp: L x 1, DataFrame
+ """
+ data_stamp = pd.DataFrame(data_stamp)
+ assert len(data.shape) == 2, 'expect inference data to have two dimensions'
+ assert len(data_stamp.shape) == 2, 'expect inference timestamps to be shaped as [L, 1]'
+ time_deltas = data_stamp.diff().dropna()
+ mean_timedelta = time_deltas.mean()[0]
+ data = data.values
+ if data.shape[0] < self.model_input_len:
+ extra_len = self.model_input_len - data.shape[0]
+ data = np.concatenate([np.mean(data, axis=0, keepdims=True).repeat(extra_len, axis=0), data], axis=0)
+ extrapolated_timestamp = pd.date_range(data_stamp[0][0] - extra_len * mean_timedelta, periods=extra_len,
+ freq=mean_timedelta)
+ data_stamp = pd.concat([extrapolated_timestamp.to_frame(), data_stamp])
+ else:
+ data = data[-self.model_input_len:, :]
+ data_stamp = data_stamp[-self.model_input_len:]
+ data = data[None, :] # add batch dim
+ return data, data_stamp
+
+ def generate_future_mark(self, time_stamp: pd.DataFrame, future_len: int) -> pd.DatetimeIndex:
+ time_deltas = time_stamp.diff().dropna()
+ mean_timedelta = time_deltas.mean()[0]
+ extrapolated_timestamp = pd.date_range(time_stamp.values[-1][0] + mean_timedelta, periods=future_len,
+ freq=mean_timedelta)
+ return extrapolated_timestamp[:, None]
diff --git a/iotdb-core/mlnode/iotdb/mlnode/process/trial.py b/iotdb-core/mlnode/iotdb/mlnode/process/trial.py
new file mode 100644
index 00000000000..9a96a526a95
--- /dev/null
+++ b/iotdb-core/mlnode/iotdb/mlnode/process/trial.py
@@ -0,0 +1,224 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+import time
+from abc import abstractmethod
+from typing import Dict, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+from iotdb.mlnode.algorithm.factory import create_forecast_model
+from iotdb.mlnode.algorithm.hyperparameter import HyperparameterName
+from iotdb.mlnode.algorithm.metric import build_metrics, forecast_metric_names
+from iotdb.mlnode.client import client_manager
+from iotdb.mlnode.constant import OptionsKey, ModelInputName
+from iotdb.mlnode.dataset.dataset import TsForecastDataset
+from iotdb.mlnode.log import logger
+from iotdb.mlnode.parser import ForecastTaskOptions
+from iotdb.mlnode.storage import model_storage
+from iotdb.mlnode.util import pack_input_dict
+from iotdb.thrift.common.ttypes import TrainingState
+
+
+class BasicTrial(object):
+ def __init__(
+ self,
+ model_id: str,
+ model: nn.Module,
+ task_configs: Dict,
+ dataset: Dataset
+ ):
+ self.model_id = model_id
+ self.dataset = dataset
+ self.trial_configs = task_configs
+ self.batch_size = task_configs[HyperparameterName.BATCH_SIZE.name()]
+ self.learning_rate = task_configs[HyperparameterName.LEARNING_RATE.name()]
+ self.epochs = task_configs[HyperparameterName.EPOCHS.name()]
+ self.num_workers = task_configs[HyperparameterName.NUM_WORKERS.name()]
+ self.use_gpu = task_configs[HyperparameterName.USE_GPU.name()]
+ self.device = self.__acquire_device()
+ self.model = model.to(self.device)
+
+ def __acquire_device(self):
+ if self.use_gpu:
+ raise NotImplementedError
+ else:
+ device = torch.device('cpu')
+ return device
+
+ def _build_dataloader(self) -> DataLoader:
+ """
+ Returns:
+ training dataloader built with the dataset
+ """
+ return DataLoader(
+ self.dataset,
+ shuffle=True,
+ drop_last=True,
+ batch_size=self.batch_size,
+ num_workers=self.num_workers
+ )
+
+ @abstractmethod
+ def start(self):
+ raise NotImplementedError
+
+
+class ForecastingTrainingTrial(BasicTrial):
+ def __init__(
+ self,
+ trial_id: str,
+ model_id: str,
+ task_options: ForecastTaskOptions,
+ model_hyperparameters: Dict,
+ task_hyperparameters: Dict,
+ dataset: TsForecastDataset
+ ):
+ """
+ A training trial, accept all parameters needed and train a single model.
+ """
+ self.model_configs = model_hyperparameters
+ # model_config contains the parameters which are needed to be saved after training.
+ self.model_configs[HyperparameterName.INPUT_VARS.value] = dataset.get_variable_num()
+ self.model_configs[OptionsKey.INPUT_LENGTH.value] = task_options.input_length
+ self.model_configs[OptionsKey.PREDICT_LENGTH.value] = task_options.predict_length
+ model = create_forecast_model(task_options.model_type, self.model_configs)
+
+ super(ForecastingTrainingTrial, self).__init__(model_id, model, task_hyperparameters, dataset)
+ self.trial_id = trial_id
+ self.pred_len = task_options.predict_length
+ self.dataloader = self._build_dataloader()
+ self.datanode_client = client_manager.borrow_data_node_client()
+ self.configNode_client = client_manager.borrow_config_node_client()
+ self.criterion = torch.nn.MSELoss()
+ self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
+ self.metrics_dict = build_metrics(forecast_metric_names)
+
+ def train(self, epoch: int) -> float:
+ self.model.train()
+ train_loss = []
+ epoch_time = time.time()
+
+ for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(self.dataloader):
+ self.optimizer.zero_grad()
+ batch_x = batch_x.float().to(self.device)
+ batch_y = batch_y.float().to(self.device)
+
+ batch_x_mark = batch_x_mark.float().to(self.device)
+ batch_y_mark = batch_y_mark.float().to(self.device)
+
+ # decoder input
+ dec_inp = torch.zeros_like(batch_y[:, -self.pred_len:, :]).float()
+
+ inputs = pack_input_dict(
+ batch_x=batch_x,
+ batch_x_mark=batch_x_mark,
+ dec_inp=dec_inp,
+ batch_y_mark=batch_y_mark)
+ outputs = self.model(inputs)
+
+ outputs = outputs[:, -self.pred_len:]
+ batch_y = batch_y[:, -self.pred_len:]
+
+ loss = self.criterion(outputs, batch_y)
+ train_loss.append(loss.item())
+
+ if (i + 1) % 500 == 0:
+ logger.info('\titers: {0}, epoch: {1} | loss: {2:.7f}'
+ .format(i + 1, epoch + 1, loss.item()))
+
+ loss.backward()
+ self.optimizer.step()
+
+ train_loss = np.average(train_loss)
+ logger.info('Epoch: {0} cost time: {1} | Train Loss: {2:.7f}'
+ .format(epoch + 1, time.time() - epoch_time, train_loss))
+ return train_loss
+
+ def vali(self, epoch: int) -> Tuple[float, Dict]:
+ self.model.eval()
+ val_loss = []
+ metrics_value_dict = {name: [] for name in forecast_metric_names}
+ for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(self.dataloader):
+ batch_x = batch_x.float().to(self.device)
+ batch_y = batch_y.float().to(self.device)
+
+ batch_x_mark = batch_x_mark.float().to(self.device)
+ batch_y_mark = batch_y_mark.float().to(self.device)
+
+ # decoder input
+ dec_inp = torch.zeros_like(batch_y[:, -self.pred_len:, :]).float()
+
+ inputs = pack_input_dict(batch_x, batch_x_mark, dec_inp, batch_y_mark)
+ outputs = self.model(inputs)
+
+ outputs = outputs[:, -self.pred_len:]
+ batch_y = batch_y[:, -self.pred_len:]
+
+ loss = self.criterion(outputs, batch_y)
+ val_loss.append(loss.item())
+
+ for name in forecast_metric_names:
+ metric = self.metrics_dict[name]
+ value = metric(outputs.detach().cpu().numpy(), batch_y.detach().cpu().numpy())
+ metrics_value_dict[name].append(value)
+
+ for name, value_list in metrics_value_dict.items():
+ metrics_value_dict[name] = np.average(value_list)
+
+ self.datanode_client.record_model_metrics(
+ model_id=self.model_id,
+ trial_id=self.trial_id,
+ metrics=list(metrics_value_dict.keys()),
+ values=list(metrics_value_dict.values())
+ )
+ val_loss = np.average(val_loss)
+ logger.info('Epoch: {0} Vali Loss: {1:.7f}'.format(epoch + 1, val_loss))
+ return val_loss, metrics_value_dict
+
+ def start(self) -> float:
+ """
+ Start training with the specified parameters, save the best model and report metrics to the db.
+ """
+ try:
+ best_loss = np.inf
+ best_metrics_dict = None
+ model_path = None
+ for epoch in range(self.epochs):
+ self.train(epoch)
+ val_loss, metrics_dict = self.vali(epoch)
+ if val_loss < best_loss:
+ best_loss = val_loss
+ best_metrics_dict = metrics_dict
+ model_path = model_storage.save_model(self.model,
+ self.model_configs,
+ model_id=self.model_id,
+ trial_id=self.trial_id)
+
+ logger.info(f'Trial: ({self.model_id}_{self.trial_id}) - Finished with best model saved successfully')
+ model_info = {}
+ model_info.update(best_metrics_dict)
+ model_info.update(self.trial_configs)
+ model_info['model_path'] = model_path
+ self.configNode_client.update_model_info(self.model_id, self.trial_id, model_info)
+ return best_loss
+ except Exception as e:
+ logger.warn(e)
+ raise e
diff --git a/iotdb-core/mlnode/iotdb/mlnode/serde.py b/iotdb-core/mlnode/iotdb/mlnode/serde.py
index 26860faf386..5b6e352866d 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/serde.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/serde.py
@@ -15,15 +15,86 @@
# specific language governing permissions and limitations
# under the License.
#
+from enum import Enum
+
import numpy as np
import pandas as pd
-from iotdb.utils.IoTDBConstants import TSDataType
+
+class TSDataType(Enum):
+ BOOLEAN = 0
+ INT32 = 1
+ INT64 = 2
+ FLOAT = 3
+ DOUBLE = 4
+ TEXT = 5
+
+ # this method is implemented to avoid the issue reported by:
+ # https://bugs.python.org/issue30545
+ def __eq__(self, other) -> bool:
+ return self.value == other.value
+
+ def __hash__(self):
+ return self.value
+
+ def np_dtype(self):
+ return {
+ TSDataType.BOOLEAN: np.dtype(">?"),
+ TSDataType.FLOAT: np.dtype(">f4"),
+ TSDataType.DOUBLE: np.dtype(">f8"),
+ TSDataType.INT32: np.dtype(">i4"),
+ TSDataType.INT64: np.dtype(">i8"),
+ TSDataType.TEXT: np.dtype("str"),
+ }[self]
+
TIMESTAMP_STR = "Time"
START_INDEX = 2
+# convert dataFrame to tsBlock in binary
+def convert_to_binary(data_frame: pd.DataFrame):
+ data_shape = data_frame.shape
+ value_column_size = data_shape[1] - 1
+ position_count = data_shape[0]
+ keys = data_frame.keys()
+
+ binary = value_column_size.to_bytes(4, byteorder="big")
+
+ # all the tsDataType are double
+ for i in range(value_column_size):
+ binary += b'\x04'
+
+ # position count
+ binary += position_count.to_bytes(4, byteorder="big")
+
+ # column encoding
+ binary += b'\x02'
+ for i in range(value_column_size):
+ binary += b'\x02'
+
+ # write columns, the column in index 0 must be timeColumn
+ binary += bool.to_bytes(False, 1, byteorder="big")
+ for i in range(position_count):
+ value = data_frame[keys[0]][i]
+ if value.dtype.byteorder != '>':
+ value = value.byteswap()
+ binary += value.tobytes()
+
+ for i in range(value_column_size):
+ # the value can't be null
+ binary += bool.to_bytes(False, 1, byteorder="big")
+ col = data_frame[keys[i + 1]]
+ for j in range(position_count):
+ value = col[j]
+ if value.dtype.byteorder != '>':
+ value = value.byteswap()
+ binary += value.tobytes()
+
+ return binary
+
+
+# convert tsBlock in binary to dataFrame
def convert_to_df(name_list, type_list, name_index, binary_list):
column_name_list = [TIMESTAMP_STR]
column_type_list = [TSDataType.INT64]
@@ -64,7 +135,7 @@ def convert_to_df(name_list, type_list, name_index, binary_list):
while binary_index < binary_size:
buffer = binary_list[binary_index]
binary_index += 1
- time_column_values, column_values, null_indicators, position_count = deserialize(buffer)
+ time_column_values, column_values, null_indicators, _ = deserialize(buffer)
time_array = np.frombuffer(
time_column_values, np.dtype(np.longlong).newbyteorder(">")
)
@@ -135,7 +206,7 @@ def convert_to_df(name_list, type_list, name_index, binary_list):
elif data_type == TSDataType.BOOLEAN:
tmp_array = np.full(total_length, np.nan, np.float32)
elif data_type == TSDataType.TEXT:
- tmp_array = np.full(total_length, None, dtype=data_array.dtype)
+ tmp_array = np.full(total_length, np.nan, dtype=data_array.dtype)
else:
raise Exception("Unsupported dataType in deserialization")
@@ -181,7 +252,7 @@ def convert_to_df(name_list, type_list, name_index, binary_list):
return df
-# Serialized tsblock:
+# Serialized tsBlock:
# +-------------+---------------+---------+------------+-----------+----------+
# | val col cnt | val col types | pos cnt | encodings | time col | val col |
# +-------------+---------------+---------+------------+-----------+----------+
@@ -199,10 +270,10 @@ def deserialize(buffer):
column_values = [None] * value_column_count
null_indicators = [None] * value_column_count
for i in range(value_column_count):
- column_value, nullIndicator, buffer = read_column(column_encodings[i + 1], buffer, data_types[i],
- position_count)
+ column_value, null_indicator, buffer = read_column(column_encodings[i + 1], buffer, data_types[i],
+ position_count)
column_values[i] = column_value
- null_indicators[i] = nullIndicator
+ null_indicators[i] = null_indicator
return time_column_values, column_values, null_indicators, position_count
@@ -228,13 +299,13 @@ def read_from_buffer(buffer, size):
def read_column_types(buffer, value_column_count):
data_types = []
- for i in range(value_column_count):
+ for _ in range(value_column_count):
res, buffer = read_byte_from_buffer(buffer)
- data_types.append(get_dataType(res))
+ data_types.append(get_data_type(res))
return data_types, buffer
-def get_dataType(value):
+def get_data_type(value):
if value == b'\x00':
return TSDataType.BOOLEAN
elif value == b'\x01':
@@ -247,15 +318,13 @@ def get_dataType(value):
return TSDataType.DOUBLE
elif value == b'\x05':
return TSDataType.TEXT
- elif value == b'\x06':
- return TSDataType.VECTOR
# Read ColumnEncodings
def read_column_encoding(buffer, size):
encodings = []
- for i in range(size):
+ for _ in range(size):
res, buffer = read_byte_from_buffer(buffer)
encodings.append(res)
return encodings, buffer
@@ -264,8 +333,8 @@ def read_column_encoding(buffer, size):
# Read Column
def deserialize_null_indicators(buffer, size):
- mayHaveNull, buffer = read_byte_from_buffer(buffer)
- if mayHaveNull != b'\x00':
+ may_have_null, buffer = read_byte_from_buffer(buffer)
+ if may_have_null != b'\x00':
return deserialize_from_boolean_array(buffer, size)
return None, buffer
@@ -278,8 +347,8 @@ def deserialize_null_indicators(buffer, size):
# +---------------+-----------------+-------------+
def read_time_column(buffer, size):
- nullIndicators, buffer = deserialize_null_indicators(buffer, size)
- if nullIndicators is None:
+ null_indicators, buffer = deserialize_null_indicators(buffer, size)
+ if null_indicators is None:
values, buffer = read_from_buffer(
buffer, size * 8
)
@@ -288,16 +357,16 @@ def read_time_column(buffer, size):
return values, buffer
-def read_INT64_column(buffer, data_type, position_count):
- nullIndicators, buffer = deserialize_null_indicators(buffer, position_count)
- if nullIndicators is None:
+def read_int64_column(buffer, data_type, position_count):
+ null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+ if null_indicators is None:
size = position_count
else:
- size = nullIndicators.count(False)
+ size = null_indicators.count(False)
if TSDataType.INT64 == data_type or TSDataType.DOUBLE == data_type:
values, buffer = read_from_buffer(buffer, size * 8)
- return values, nullIndicators, buffer
+ return values, null_indicators, buffer
else:
raise Exception("Invalid data type: " + data_type)
@@ -309,16 +378,16 @@ def read_INT64_column(buffer, data_type, position_count):
# | byte | list[byte] | list[int32] |
# +---------------+-----------------+-------------+
-def read_Int32_column(buffer, data_type, position_count):
- nullIndicators, buffer = deserialize_null_indicators(buffer, position_count)
- if nullIndicators is None:
+def read_int32_column(buffer, data_type, position_count):
+ null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+ if null_indicators is None:
size = position_count
else:
- size = nullIndicators.count(False)
+ size = null_indicators.count(False)
if TSDataType.INT32 == data_type or TSDataType.FLOAT == data_type:
values, buffer = read_from_buffer(buffer, size * 4)
- return values, nullIndicators, buffer
+ return values, null_indicators, buffer
else:
raise Exception("Invalid data type: " + data_type)
@@ -333,9 +402,9 @@ def read_Int32_column(buffer, data_type, position_count):
def read_byte_column(buffer, data_type, position_count):
if data_type != TSDataType.BOOLEAN:
raise Exception("Invalid data type: " + data_type)
- nullIndicators, buffer = deserialize_null_indicators(buffer, position_count)
+ null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
res, buffer = deserialize_from_boolean_array(buffer, position_count)
- return res, nullIndicators, buffer
+ return res, null_indicators, buffer
def deserialize_from_boolean_array(buffer, size):
@@ -386,31 +455,31 @@ def deserialize_from_boolean_array(buffer, size):
def read_binary_column(buffer, data_type, position_count):
if data_type != TSDataType.TEXT:
raise Exception("Invalid data type: " + data_type)
- nullIndicators, buffer = deserialize_null_indicators(buffer, position_count)
+ null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
- if nullIndicators is None:
+ if null_indicators is None:
size = position_count
else:
- size = nullIndicators.count(False)
+ size = null_indicators.count(False)
values = [None] * size
for i in range(size):
length, buffer = read_int_from_buffer(buffer)
res, buffer = read_from_buffer(buffer, length)
values[i] = res
- return values, nullIndicators, buffer
+ return values, null_indicators, buffer
def read_column(encoding, buffer, data_type, position_count):
if encoding == b'\x00':
return read_byte_column(buffer, data_type, position_count)
elif encoding == b'\x01':
- return read_Int32_column(buffer, data_type, position_count)
+ return read_int32_column(buffer, data_type, position_count)
elif encoding == b'\x02':
- return read_INT64_column(buffer, data_type, position_count)
+ return read_int64_column(buffer, data_type, position_count)
elif encoding == b'\x03':
return read_binary_column(buffer, data_type, position_count)
elif encoding == b'\x04':
- return read_runLength_column(buffer, data_type, position_count)
+ return read_run_length_column(buffer, data_type, position_count)
else:
raise Exception("Unsupported encoding: " + encoding)
@@ -422,11 +491,11 @@ def read_column(encoding, buffer, data_type, position_count):
# | byte | list[byte] |
# +-----------+-------------------------+
-def read_runLength_column(buffer, data_type, position_count):
+def read_run_length_column(buffer, data_type, position_count):
encoding, buffer = read_byte_from_buffer(buffer)
- column, nullIndicators, buffer = read_column(encoding, buffer, data_type, 1)
+ column, null_indicators, buffer = read_column(encoding, buffer, data_type, 1)
- return repeat(column, data_type, position_count), nullIndicators * position_count, buffer
+ return repeat(column, data_type, position_count), null_indicators * position_count, buffer
def repeat(buffer, data_type, position_count):
diff --git a/iotdb-core/mlnode/iotdb/mlnode/service.py b/iotdb-core/mlnode/iotdb/mlnode/service.py
index 8bf97402d1f..30c963aa304 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/service.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/service.py
@@ -56,7 +56,3 @@ class MLNode(object):
time.sleep(0.1)
logger.info('IoTDB-MLNode has successfully started.')
-
-if __name__ == "__main__":
- server = MLNode()
- server.start()
diff --git a/iotdb-core/mlnode/iotdb/mlnode/storage.py b/iotdb-core/mlnode/iotdb/mlnode/storage.py
index 84d7dfd7ed1..8f79ecd64c3 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/storage.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/storage.py
@@ -19,6 +19,7 @@
import json
import os
import shutil
+import threading
from typing import Dict, Tuple
import torch
@@ -26,20 +27,22 @@ import torch.nn as nn
from pylru import lrucache
from iotdb.mlnode.config import descriptor
+from iotdb.mlnode.constant import OptionsKey
from iotdb.mlnode.exception import ModelNotExistError
from iotdb.mlnode.log import logger
+from iotdb.mlnode.util import pack_input_dict
class ModelStorage(object):
def __init__(self):
- self.__model_dir = os.path.join('.', descriptor.get_config().get_mn_model_storage_dir())
+ self.__model_dir = os.path.join(os.getcwd(), descriptor.get_config().get_mn_model_storage_dir())
if not os.path.exists(self.__model_dir):
try:
os.mkdir(self.__model_dir)
except PermissionError as e:
logger.error(e)
raise e
-
+ self.lock = threading.RLock()
self.__model_cache = lrucache(descriptor.get_config().get_mn_model_storage_cache_size())
def save_model(self,
@@ -47,28 +50,34 @@ class ModelStorage(object):
model_config: Dict,
model_id: str,
trial_id: str) -> str:
- """
- Note: model config for time series should contain 'input_len' and 'input_vars'
- """
model_dir_path = os.path.join(self.__model_dir, f'{model_id}')
+ self.lock.acquire()
if not os.path.exists(model_dir_path):
os.makedirs(model_dir_path)
+ self.lock.release()
model_file_path = os.path.join(model_dir_path, f'{trial_id}.pt')
- sample_input = [torch.randn(1, model_config['input_len'], model_config['input_vars'])]
+ # Note: model config for time series should contain 'input_len' and 'input_vars'
+ sample_input = (
+ pack_input_dict(
+ torch.randn(1, model_config[OptionsKey.INPUT_LENGTH.name()], model_config[OptionsKey.INPUT_VARS.name()])
+ )
+ )
+ self.lock.acquire()
torch.jit.save(torch.jit.trace(model, sample_input),
model_file_path,
_extra_files={'model_config': json.dumps(model_config)})
+ self.lock.release()
return os.path.abspath(model_file_path)
- def load_model(self, model_id: str, trial_id: str) -> Tuple[torch.jit.ScriptModule, Dict]:
+ def load_model(self, file_path: str) -> Tuple[torch.jit.ScriptModule, Dict]:
"""
Returns:
jit_model: a ScriptModule contains model architecture and parameters, which can be deployed cross-platform
model_config: a dict contains model attributes
"""
- file_path = os.path.join(self.__model_dir, f'{model_id}', f'{trial_id}.pt')
- if model_id in self.__model_cache:
+ file_path = os.path.join(self.__model_dir, file_path)
+ if file_path in self.__model_cache:
return self.__model_cache[file_path]
else:
if not os.path.exists(file_path):
@@ -98,5 +107,4 @@ class ModelStorage(object):
del self.__model_cache[file_path]
-# initialize a singleton
model_storage = ModelStorage()
diff --git a/iotdb-core/mlnode/iotdb/mlnode/util.py b/iotdb-core/mlnode/iotdb/mlnode/util.py
index 5cdc52f01af..abb1c9c55fe 100644
--- a/iotdb-core/mlnode/iotdb/mlnode/util.py
+++ b/iotdb-core/mlnode/iotdb/mlnode/util.py
@@ -15,7 +15,11 @@
# specific language governing permissions and limitations
# under the License.
#
-from iotdb.mlnode.constant import TSStatusCode
+import os
+
+import torch
+
+from iotdb.mlnode.constant import TSStatusCode, ModelInputName
from iotdb.mlnode.exception import BadNodeUrlError
from iotdb.mlnode.log import logger
from iotdb.thrift.common.ttypes import TEndPoint, TSStatus
@@ -55,3 +59,22 @@ def verify_success(status: TSStatus, err_msg: str) -> None:
if status.code != TSStatusCode.SUCCESS_STATUS.get_status_code():
logger.warn(err_msg + ", error status is ", status)
raise RuntimeError(str(status.code) + ": " + status.message)
+
+
+def pack_input_dict(batch_x: torch.Tensor,
+ batch_x_mark: torch.Tensor = None,
+ dec_inp: torch.Tensor = None,
+ batch_y_mark: torch.Tensor = None):
+ """
+ pack up inputs as a dict to adapt for different models
+ """
+ input_dict = {}
+ if batch_x is not None:
+ input_dict[ModelInputName.DATA_X.value] = batch_x
+ if batch_x_mark is not None:
+ input_dict[ModelInputName.TIME_STAMP_X] = batch_x_mark
+ if dec_inp is not None:
+ input_dict[ModelInputName.DEC_INP] = dec_inp
+ if batch_y_mark is not None:
+ input_dict[ModelInputName.TIME_STAMP_Y.value] = batch_y_mark
+ return input_dict
diff --git a/iotdb-core/mlnode/requirements.txt b/iotdb-core/mlnode/requirements.txt
index c49c8a01896..0077b856855 100644
--- a/iotdb-core/mlnode/requirements.txt
+++ b/iotdb-core/mlnode/requirements.txt
@@ -21,7 +21,10 @@ numpy>=1.21.4
apache-iotdb
poetry
torch~=2.0.0
-pylru
+pylru~=1.2.1
thrift~=0.13.0
-dynaconf~=3.1.12
\ No newline at end of file
+dynaconf~=3.1.12
+requests~=2.31.0
+optuna~=3.2.0
+psutil~=5.9.5
\ No newline at end of file
diff --git a/iotdb-core/mlnode/resources/conf/iotdb-mlnode.toml b/iotdb-core/mlnode/resources/conf/iotdb-mlnode.toml
index a029509e643..474e86bfadb 100644
--- a/iotdb-core/mlnode/resources/conf/iotdb-mlnode.toml
+++ b/iotdb-core/mlnode/resources/conf/iotdb-mlnode.toml
@@ -37,7 +37,19 @@ mn_model_storage_dir = "models"
# Cache number of model storage to avoid repeated loading
# Datatype: int
-mn_model_storage_cachesize = 30
+mn_model_storage_cache_size = 30
+
+# Maximum number of training model tasks, otherwise the task is pending
+# Datatype: int
+mn_task_pool_size = 10
+
+# Maximum number of trials to be explored in a tuning task
+# Datatype: int
+mn_tuning_trial_num = 20
+
+# Concurrency of trials in a tuning task
+# Datatype: int
+mn_tuning_trial_concurrency = 4
####################
### Target Config Node
diff --git a/iotdb-core/mlnode/test/test_create_forecast_model.py b/iotdb-core/mlnode/test/test_create_forecast_model.py
new file mode 100644
index 00000000000..7a13fc6f6f4
--- /dev/null
+++ b/iotdb-core/mlnode/test/test_create_forecast_model.py
@@ -0,0 +1,100 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+import torch
+from iotdb.mlnode.util import pack_input_dict
+
+from iotdb.mlnode.algorithm.factory import (ForecastModelType,
+ create_forecast_model)
+from iotdb.mlnode.algorithm.hyperparameter import (HyperparameterName,
+ parse_fixed_hyperparameters)
+from iotdb.mlnode.constant import OptionsKey
+from iotdb.mlnode.exception import (BadConfigValueError, MissingOptionError,
+ UnsupportedError)
+from iotdb.mlnode.parser import ForecastTaskOptions, parse_task_options
+
+
+def test_create_forecast_model():
+ d_forecast_task_options = ForecastTaskOptions({OptionsKey.MODEL_TYPE.name(): ForecastModelType.DLINEAR.value,
+ OptionsKey.AUTO_TUNING.name(): "False"})
+ d_configs = {HyperparameterName.KERNEL_SIZE.name(): "25", HyperparameterName.USE_GPU.name(): "False"}
+ model_configs, _ = parse_fixed_hyperparameters(d_forecast_task_options, d_configs)
+ model_configs[HyperparameterName.INPUT_VARS.value] = 8
+ model_configs[OptionsKey.INPUT_LENGTH.value] = d_forecast_task_options.input_length
+ model_configs[OptionsKey.PREDICT_LENGTH.value] = d_forecast_task_options.predict_length
+ model = create_forecast_model(d_forecast_task_options.model_type, model_configs)
+ sample_input = torch.randn(1, model_configs[OptionsKey.INPUT_LENGTH.value],
+ model_configs[HyperparameterName.INPUT_VARS.value])
+ sample_input = pack_input_dict(sample_input)
+ output = model(sample_input)
+ assert output.shape[1] == model_configs[OptionsKey.PREDICT_LENGTH.value]
+ assert output.shape[2] == model_configs[HyperparameterName.INPUT_VARS.value]
+ assert model_configs[HyperparameterName.KERNEL_SIZE.value] == 25
+
+ n_forecast_task_options = ForecastTaskOptions({OptionsKey.MODEL_TYPE.name(): ForecastModelType.NBEATS.value})
+ n_configs = {HyperparameterName.KERNEL_SIZE.name(): "25", HyperparameterName.D_MODEL.name(): "64"}
+ model_configs, _ = parse_fixed_hyperparameters(n_forecast_task_options, n_configs)
+ model_configs[HyperparameterName.INPUT_VARS.name()] = 8
+ model_configs[OptionsKey.INPUT_LENGTH.name()] = n_forecast_task_options.input_length
+ model_configs[OptionsKey.PREDICT_LENGTH.name()] = n_forecast_task_options.predict_length
+ model = create_forecast_model(n_forecast_task_options.model_type, model_configs)
+ output = model(sample_input)
+ assert output.shape[1] == model_configs[OptionsKey.PREDICT_LENGTH.value]
+ assert output.shape[2] == model_configs[HyperparameterName.INPUT_VARS.value]
+ assert model_configs['d_model'] == 64
+ assert model_configs['block_type'] == "generic"
+ assert 'kernel_size' not in model_configs # config kernel_size not belongs to nbeats model
+
+
+def test_bad_config_model1():
+ try:
+ d_forecast_task_options = ForecastTaskOptions({OptionsKey.MODEL_TYPE.name(): "dlinear_dummy",
+ OptionsKey.AUTO_TUNING.name(): "False"})
+ d_configs = {HyperparameterName.KERNEL_SIZE.name(): "25", HyperparameterName.USE_GPU.name(): "False"}
+ create_forecast_model(d_forecast_task_options.model_type, d_configs)
+ except UnsupportedError as e:
+ assert e.message == "model_type dlinear_dummy is not supported in current version"
+
+
+def test_bad_config_model2():
+ try:
+ d_forecast_task_options = ForecastTaskOptions({OptionsKey.MODEL_TYPE.name(): "dlinear",
+ OptionsKey.AUTO_TUNING.name(): "False"})
+ d_configs = {HyperparameterName.KERNEL_SIZE.name(): "-1", HyperparameterName.USE_GPU.name(): "False"}
+ parse_fixed_hyperparameters(d_forecast_task_options, d_configs)
+ except BadConfigValueError as e:
+ assert e.message == "Bad value [-1] for config kernel_size. Expect value between 1 and 10000000000.0, " \
+ "got -1 instead."
+
+
+def test_bad_config_model3():
+ try:
+ d_task_options = {OptionsKey.MODEL_TYPE.name(): "dlinear",
+ OptionsKey.AUTO_TUNING.name(): False}
+ parse_task_options(d_task_options)
+ except MissingOptionError as e:
+ assert e.message == "Missing task option: task_type"
+
+
+def test_bad_config_model5():
+ try:
+ d_task_options = {OptionsKey.MODEL_TYPE.name(): "dlinear",
+ OptionsKey.TASK_TYPE.name(): "dummy_task",
+ OptionsKey.AUTO_TUNING.name(): False}
+ parse_task_options(d_task_options)
+ except UnsupportedError as e:
+ assert e.message == "task_type dummy_task is not supported in current version"
diff --git a/iotdb-core/mlnode/test/test_model_storage.py b/iotdb-core/mlnode/test/test_model_storage.py
index 863e73b7164..6eea0446d92 100644
--- a/iotdb-core/mlnode/test/test_model_storage.py
+++ b/iotdb-core/mlnode/test/test_model_storage.py
@@ -18,8 +18,10 @@
import os
import time
+from typing import Dict
import torch.nn as nn
+from iotdb.mlnode.constant import ModelInputName
from iotdb.mlnode.config import descriptor
from iotdb.mlnode.exception import ModelNotExistError
@@ -31,13 +33,15 @@ class ExampleModel(nn.Module):
super(ExampleModel, self).__init__()
self.layer = nn.Identity()
- def forward(self, x):
+ def forward(self, input_dict: Dict):
+ # x: [Batch, Input length, Channel]
+ x = input_dict[ModelInputName.DATA_X.value]
return self.layer(x)
model = ExampleModel()
model_config = {
- 'input_len': 1,
+ 'input_length': 1,
'input_vars': 1,
'id': time.time()
}
@@ -54,8 +58,11 @@ def test_save_model():
def test_load_model():
trial_id = 'tid_0'
model_id = 'mid_test_model_load'
+
+ model_file_path = os.path.join(f'{model_id}', f'{trial_id}.pt')
+
model_storage.save_model(model, model_config, model_id=model_id, trial_id=trial_id)
- model_loaded, model_config_loaded = model_storage.load_model(model_id=model_id, trial_id=trial_id)
+ _, model_config_loaded = model_storage.load_model(model_file_path)
assert model_config == model_config_loaded
@@ -63,10 +70,10 @@ def test_load_not_exist_model():
trial_id = 'dummy_trial'
model_id = 'dummy_model'
try:
- model_loaded, model_config_loaded = model_storage.load_model(model_id=model_id, trial_id=trial_id)
+ model_storage.load_model(os.path.join(f'{model_id}', f'{trial_id}.pt'))
except Exception as e:
assert e.message == ModelNotExistError(
- os.path.join('.', descriptor.get_config().get_mn_model_storage_dir(),
+ os.path.join(os.getcwd(), descriptor.get_config().get_mn_model_storage_dir(),
model_id, f'{trial_id}.pt')).message
diff --git a/iotdb-core/mlnode/test/test_serde.py b/iotdb-core/mlnode/test/test_serde.py
index c05083be417..f4a3052782c 100644
--- a/iotdb-core/mlnode/test/test_serde.py
+++ b/iotdb-core/mlnode/test/test_serde.py
@@ -26,7 +26,7 @@ device_id = "root.wt1"
ts_path_lst = [
"root.wt1.temperature",
- "root.wt1.windspeed",
+ "root.wt1.windSpeed",
"root.wt1.angle",
"root.wt1.altitude",
"root.wt1.status",
@@ -34,7 +34,7 @@ ts_path_lst = [
]
measurements = [
"temperature",
- "windspeed",
+ "windSpeed",
"angle",
"altitude",
"status",
@@ -90,10 +90,10 @@ binary_with_null = \
b'\x00\x00\x00\x06\x04\x03\x05\x00\x02\x01\x00\x00\x00\x01\x02\x02\x01\x03\x04\x02\x04\x00\x00\x00\x00\x00'
b'\x00\x00\x00\x13\x00?\xd3\xc7\xae#\x86\xe1j\x00?\x10\xea!\x00\x00\x00\x00\x05text1\x00\x01\x80\x00\x00\x00'
b'\x00\x00\x00\x00\x00\x00?\x01\x01\x80']
-column_names = ['root.wt1.altitude', 'root.wt1.temperature', 'root.wt1.angle', 'root.wt1.windspeed',
+column_names = ['root.wt1.altitude', 'root.wt1.temperature', 'root.wt1.angle', 'root.wt1.windSpeed',
'root.wt1.hardware', 'root.wt1.status']
data_type_list = ['INT64', 'FLOAT', 'INT32', 'DOUBLE', 'TEXT', 'BOOLEAN']
-column_name_index = {'root.wt1.windspeed': 0, 'root.wt1.temperature': 1, 'root.wt1.hardware': 2, 'root.wt1.status': 3,
+column_name_index = {'root.wt1.windSpeed': 0, 'root.wt1.temperature': 1, 'root.wt1.hardware': 2, 'root.wt1.status': 3,
'root.wt1.altitude': 4, 'root.wt1.angle': 5}
diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-common.properties b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-common.properties
index a8c1773ac7a..277e7c3d17c 100644
--- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-common.properties
+++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-common.properties
@@ -1121,3 +1121,9 @@ cluster_name=defaultCluster
# SSL timeout (in seconds)
# idle_timeout_in_seconds=50000
+####################
+### IoTDB-ML Configuration
+####################
+
+# The thread count which can be used for model inference operation.
+# model_inference_execution_thread_count=5
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ClientPoolFactory.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ClientPoolFactory.java
index 9297136365f..97882f0c963 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ClientPoolFactory.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ClientPoolFactory.java
@@ -24,6 +24,7 @@ import org.apache.iotdb.commons.client.async.AsyncConfigNodeIServiceClient;
import org.apache.iotdb.commons.client.async.AsyncDataNodeInternalServiceClient;
import org.apache.iotdb.commons.client.async.AsyncDataNodeMPPDataExchangeServiceClient;
import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient;
+import org.apache.iotdb.commons.client.mlnode.MLNodeClient;
import org.apache.iotdb.commons.client.property.ClientPoolProperty;
import org.apache.iotdb.commons.client.property.ThriftClientProperty;
import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
@@ -282,4 +283,29 @@ public class ClientPoolFactory {
return clientPool;
}
}
+
+ public static class MLNodeClientPoolFactory
+ implements IClientPoolFactory<TEndPoint, MLNodeClient> {
+
+ @Override
+ public KeyedObjectPool<TEndPoint, MLNodeClient> createClientPool(
+ ClientManager<TEndPoint, MLNodeClient> manager) {
+ GenericKeyedObjectPool<TEndPoint, MLNodeClient> clientPool =
+ new GenericKeyedObjectPool<>(
+ new MLNodeClient.Factory(
+ manager,
+ new ThriftClientProperty.Builder()
+ .setConnectionTimeoutMs(conf.getConnectionTimeoutInMS())
+ .setRpcThriftCompressionEnabled(true)
+ .build()),
+ new ClientPoolProperty.Builder<MLNodeClient>()
+ .setCoreClientNumForEachNode(conf.getCoreClientNumForEachNode())
+ .setMaxClientNumForEachNode(conf.getMaxClientNumForEachNode())
+ .build()
+ .getConfig());
+ ClientManagerMetrics.getInstance()
+ .registerClientManager(this.getClass().getSimpleName(), clientPool);
+ return clientPool;
+ }
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/MLNodeClient.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClient.java
similarity index 50%
rename from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/MLNodeClient.java
rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClient.java
index 8454e536ff1..a925791498f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/MLNodeClient.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClient.java
@@ -17,12 +17,14 @@
* under the License.
*/
-package org.apache.iotdb.db.protocol.client;
+package org.apache.iotdb.commons.client.mlnode;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
-import org.apache.iotdb.commons.client.property.ClientPoolProperty;
-import org.apache.iotdb.commons.conf.CommonDescriptor;
+import org.apache.iotdb.commons.client.ClientManager;
+import org.apache.iotdb.commons.client.ThriftClient;
+import org.apache.iotdb.commons.client.factory.ThriftClientFactory;
+import org.apache.iotdb.commons.client.property.ThriftClientProperty;
import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.mlnode.rpc.thrift.IMLNodeRPCService;
import org.apache.iotdb.mlnode.rpc.thrift.TCreateTrainingTaskReq;
@@ -30,13 +32,13 @@ import org.apache.iotdb.mlnode.rpc.thrift.TDeleteModelReq;
import org.apache.iotdb.mlnode.rpc.thrift.TForecastReq;
import org.apache.iotdb.mlnode.rpc.thrift.TForecastResp;
import org.apache.iotdb.rpc.TConfigurationConst;
-import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.read.common.block.TsBlock;
import org.apache.iotdb.tsfile.read.common.block.column.TsBlockSerde;
+import org.apache.commons.pool2.PooledObject;
+import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.thrift.TException;
-import org.apache.thrift.protocol.TCompactProtocol;
-import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
@@ -45,61 +47,75 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
import java.util.Map;
import java.util.Optional;
-public class MLNodeClient implements AutoCloseable {
+public class MLNodeClient implements AutoCloseable, ThriftClient {
private static final Logger logger = LoggerFactory.getLogger(MLNodeClient.class);
- private final TTransport transport;
- private final IMLNodeRPCService.Client client;
+ private static final TEndPoint endPoint = MLNodeInfo.endPoint;
+
+ private TTransport transport;
+
+ private final ThriftClientProperty property;
+ private IMLNodeRPCService.Client client;
public static final String MSG_CONNECTION_FAIL =
"Fail to connect to MLNode. Please check status of MLNode";
private final TsBlockSerde tsBlockSerde = new TsBlockSerde();
- public MLNodeClient() throws TException {
- TEndPoint endpoint = CommonDescriptor.getInstance().getConfig().getTargetMLNodeEndPoint();
+ ClientManager<TEndPoint, MLNodeClient> clientManager;
+
+ public MLNodeClient(
+ ThriftClientProperty property, ClientManager<TEndPoint, MLNodeClient> clientManager)
+ throws TException {
+ this.property = property;
+ this.clientManager = clientManager;
+ init();
+ }
+
+ private void init() throws TException {
try {
- long connectionTimeout = ClientPoolProperty.DefaultProperty.WAIT_CLIENT_TIMEOUT_MS;
transport =
new TFramedTransport.Factory()
.getTransport(
new TSocket(
TConfigurationConst.defaultTConfiguration,
- endpoint.getIp(),
- endpoint.getPort(),
- (int) connectionTimeout));
+ endPoint.getIp(),
+ endPoint.getPort(),
+ property.getConnectionTimeoutMs()));
if (!transport.isOpen()) {
transport.open();
}
} catch (TTransportException e) {
throw new TException(MSG_CONNECTION_FAIL);
}
+ client = new IMLNodeRPCService.Client(property.getProtocolFactory().getProtocol(transport));
+ }
- TProtocolFactory protocolFactory = new TCompactProtocol.Factory();
- client = new IMLNodeRPCService.Client(protocolFactory.getProtocol(transport));
+ public TTransport getTransport() {
+ return transport;
}
public TSStatus createTrainingTask(
- ModelInformation modelInformation, Map<String, String> modelConfigs) throws TException {
+ ModelInformation modelInformation, Map<String, String> hyperparameters) throws TException {
try {
TCreateTrainingTaskReq req =
new TCreateTrainingTaskReq(
modelInformation.getModelId(),
- modelInformation.isAuto(),
- modelConfigs,
- modelInformation.getQueryExpressions());
- if (modelInformation.getQueryFilter() != null) {
- req.setQueryFilter(modelInformation.getQueryFilter());
- }
+ modelInformation.getOptions(),
+ hyperparameters,
+ modelInformation.getDatasetFetchSql());
return client.createTrainingTask(req);
} catch (TException e) {
logger.warn(
- "Failed to connect to MLNode from ConfigNode when executing {}",
- Thread.currentThread().getStackTrace()[1].getMethodName());
+ "Failed to connect to MLNode from ConfigNode when executing {}: {}",
+ Thread.currentThread().getStackTrace()[1].getMethodName(),
+ e.getMessage());
throw new TException(MSG_CONNECTION_FAIL);
}
}
@@ -109,26 +125,40 @@ public class MLNodeClient implements AutoCloseable {
return client.deleteModel(new TDeleteModelReq(modelId));
} catch (TException e) {
logger.warn(
- "Failed to connect to MLNode from ConfigNode when executing {}",
- Thread.currentThread().getStackTrace()[1].getMethodName());
+ "Failed to connect to MLNode from ConfigNode when executing {}: {}",
+ Thread.currentThread().getStackTrace()[1].getMethodName(),
+ e.getMessage());
throw new TException(MSG_CONNECTION_FAIL);
}
}
- public TsBlock forecast(String modelPath, TsBlock inputTsBlock) throws TException {
+ public TForecastResp forecast(
+ String modelPath,
+ TsBlock inputTsBlock,
+ List<TSDataType> inputTypeList,
+ List<String> inputColumnNameList,
+ int predictLength)
+ throws TException {
try {
- TForecastReq forecastReq = new TForecastReq(modelPath, tsBlockSerde.serialize(inputTsBlock));
- TForecastResp resp = client.forecast(forecastReq);
- if (resp.status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- throw new TException("Failed to execute forecast task, because: " + resp.status.message);
+ List<String> reqInputTypeList = new ArrayList<>();
+ for (TSDataType dataType : inputTypeList) {
+ reqInputTypeList.add(dataType.toString());
}
- return tsBlockSerde.deserialize(resp.forecastResult);
+ TForecastReq forecastReq =
+ new TForecastReq(
+ modelPath,
+ tsBlockSerde.serialize(inputTsBlock),
+ reqInputTypeList,
+ inputColumnNameList,
+ predictLength);
+ return client.forecast(forecastReq);
} catch (IOException e) {
throw new TException("An exception occurred while serializing input tsblock", e);
} catch (TException e) {
logger.warn(
- "Failed to connect to MLNode from DataNode when executing {}",
- Thread.currentThread().getStackTrace()[1].getMethodName());
+ "Failed to connect to MLNode from DataNode when executing {}: {}",
+ Thread.currentThread().getStackTrace()[1].getMethodName(),
+ e.getMessage());
throw new TException(MSG_CONNECTION_FAIL);
}
}
@@ -137,4 +167,46 @@ public class MLNodeClient implements AutoCloseable {
public void close() throws Exception {
Optional.ofNullable(transport).ifPresent(TTransport::close);
}
+
+ @Override
+ public void invalidate() {
+ Optional.ofNullable(transport).ifPresent(TTransport::close);
+ }
+
+ @Override
+ public void invalidateAll() {
+ clientManager.clear(endPoint);
+ }
+
+ @Override
+ public boolean printLogWhenEncounterException() {
+ return property.isPrintLogWhenEncounterException();
+ }
+
+ public static class Factory extends ThriftClientFactory<TEndPoint, MLNodeClient> {
+
+ public Factory(
+ ClientManager<TEndPoint, MLNodeClient> clientClientManager,
+ ThriftClientProperty thriftClientProperty) {
+ super(clientClientManager, thriftClientProperty);
+ }
+
+ @Override
+ public void destroyObject(TEndPoint tEndPoint, PooledObject<MLNodeClient> pooledObject)
+ throws Exception {
+ pooledObject.getObject().close();
+ }
+
+ @Override
+ public PooledObject<MLNodeClient> makeObject(TEndPoint tEndPoint) throws Exception {
+ return new DefaultPooledObject<>(new MLNodeClient(thriftClientProperty, clientManager));
+ }
+
+ @Override
+ public boolean validateObject(TEndPoint tEndPoint, PooledObject<MLNodeClient> pooledObject) {
+ return Optional.ofNullable(pooledObject.getObject().getTransport())
+ .map(TTransport::isOpen)
+ .orElse(false);
+ }
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClientManager.java
similarity index 53%
copy from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
copy to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClientManager.java
index 336bbcce0bb..25cd4a12908 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeClientManager.java
@@ -17,11 +17,24 @@
* under the License.
*/
-package org.apache.iotdb.db.queryengine.plan.expression.multi;
+package org.apache.iotdb.commons.client.mlnode;
-/** */
-public enum FunctionType {
- AGGREGATION_FUNCTION,
- BUILT_IN_SCALAR_FUNCTION,
- UDF
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
+import org.apache.iotdb.commons.client.IClientManager;
+
+public class MLNodeClientManager {
+ private MLNodeClientManager() {
+ // Empty constructor
+ }
+
+ private static final class MLNodeClientManagerHolder {
+ private static final IClientManager<TEndPoint, MLNodeClient> INSTANCE =
+ new IClientManager.Factory<TEndPoint, MLNodeClient>()
+ .createClientManager(new ClientPoolFactory.MLNodeClientPoolFactory());
+ }
+
+ public static IClientManager<TEndPoint, MLNodeClient> getInstance() {
+ return MLNodeClientManagerHolder.INSTANCE;
+ }
}
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeInfo.java
similarity index 70%
copy from iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
copy to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeInfo.java
index 304d209817f..be29febd687 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/state/model/CreateModelState.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/mlnode/MLNodeInfo.java
@@ -17,11 +17,13 @@
* under the License.
*/
-package org.apache.iotdb.confignode.procedure.state.model;
+package org.apache.iotdb.commons.client.mlnode;
-public enum CreateModelState {
- INIT,
- VALIDATED,
- CONFIG_NODE_ACTIVE,
- ML_NODE_ACTIVE
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.conf.CommonDescriptor;
+
+public class MLNodeInfo {
+ // currently, we only support one MLNode
+ public static final TEndPoint endPoint =
+ CommonDescriptor.getInstance().getConfig().getTargetMLNodeEndPoint();
}
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java
index 4226adb5b67..96dfbb0f80d 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java
@@ -137,6 +137,9 @@ public class CommonConfig {
/** This variable set timestamp precision as millisecond, microsecond or nanosecond. */
private String timestampPrecision = "ms";
+ /** The number of threads in the thread pool that execute model inference tasks. */
+ private int modelInferenceExecutionThreadCount = 5;
+
/**
* The name of the directory that stores the tsfiles temporarily hold or generated by the pipe
* module. The directory is located in the data directory of IoTDB.
@@ -752,4 +755,12 @@ public class CommonConfig {
public void setDatabaseLimitThreshold(int databaseLimitThreshold) {
this.databaseLimitThreshold = databaseLimitThreshold;
}
+
+ public int getModelInferenceExecutionThreadCount() {
+ return modelInferenceExecutionThreadCount;
+ }
+
+ public void setModelInferenceExecutionThreadCount(int modelInferenceExecutionThreadCount) {
+ this.modelInferenceExecutionThreadCount = modelInferenceExecutionThreadCount;
+ }
}
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonDescriptor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonDescriptor.java
index 0922b96452d..ba8bf0c9b04 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonDescriptor.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonDescriptor.java
@@ -238,6 +238,11 @@ public class CommonDescriptor {
Integer.parseInt(
properties.getProperty(
"database_limit_threshold", String.valueOf(config.getDatabaseLimitThreshold()))));
+ config.setModelInferenceExecutionThreadCount(
+ Integer.parseInt(
+ properties.getProperty(
+ "model_inference_execution_thread_count",
+ String.valueOf(config.getModelInferenceExecutionThreadCount()))));
}
private void loadPipeProps(Properties properties) {
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ForecastModeInformation.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ForecastModeInformation.java
new file mode 100644
index 00000000000..3b3eebc1f47
--- /dev/null
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ForecastModeInformation.java
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.model;
+
+import org.apache.iotdb.common.rpc.thrift.TaskType;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
+import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class ForecastModeInformation extends ModelInformation {
+
+ private final List<TSDataType> inputTypeList;
+
+ private final List<Integer> predictIndexList;
+
+ private final int inputLength;
+ private final int predictLength;
+
+ public static final String INPUT_TYPE_LIST = "input_type_list";
+ public static final String PREDICT_INDEX_LIST = "predict_index_list";
+ public static final String INPUT_LENGTH = "input_length";
+ public static final String PREDICT_LENGTH = "predict_length";
+
+ public static final String DEFAULT_INPUT_LENGTH = "96";
+ public static final String DEFAULT_PREDICT_LENGTH = "96";
+
+ public ForecastModeInformation(
+ String modelId,
+ Map<String, String> options,
+ String datasetFetchSQL,
+ List<TSDataType> inputTypeList,
+ List<Integer> predictIndexList,
+ int inputLength,
+ int predictLength) {
+ super(modelId, options, datasetFetchSQL);
+ this.inputTypeList = inputTypeList;
+ this.predictIndexList = predictIndexList;
+ this.inputLength = inputLength;
+ this.predictLength = predictLength;
+ }
+
+ public ForecastModeInformation(ByteBuffer buffer) {
+ super(buffer);
+ int listSize = ReadWriteIOUtils.readInt(buffer);
+ this.inputTypeList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ inputTypeList.add(TSDataType.deserializeFrom(buffer));
+ }
+ listSize = ReadWriteIOUtils.readInt(buffer);
+ this.predictIndexList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ predictIndexList.add(ReadWriteIOUtils.readInt(buffer));
+ }
+ this.inputLength = ReadWriteIOUtils.readInt(buffer);
+ this.predictLength = ReadWriteIOUtils.readInt(buffer);
+ }
+
+ public ForecastModeInformation(InputStream stream) throws IOException {
+ super(stream);
+ int listSize = ReadWriteIOUtils.readInt(stream);
+ this.inputTypeList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ inputTypeList.add(TSDataType.deserializeFrom(stream));
+ }
+ listSize = ReadWriteIOUtils.readInt(stream);
+ this.predictIndexList = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ predictIndexList.add(ReadWriteIOUtils.readInt(stream));
+ }
+ this.inputLength = ReadWriteIOUtils.readInt(stream);
+ this.predictLength = ReadWriteIOUtils.readInt(stream);
+ }
+
+ public TaskType getTaskType() {
+ return TaskType.FORECAST;
+ }
+
+ public List<TSDataType> getInputTypeList() {
+ return inputTypeList;
+ }
+
+ public List<Integer> getPredictIndexList() {
+ return predictIndexList;
+ }
+
+ public int getInputLength() {
+ return inputLength;
+ }
+
+ public int getPredictLength() {
+ return predictLength;
+ }
+
+ @Override
+ public void serialize(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(TaskType.FORECAST.ordinal(), stream);
+
+ super.serialize(stream);
+
+ ReadWriteIOUtils.write(inputTypeList.size(), stream);
+ for (TSDataType inputType : inputTypeList) {
+ inputType.serializeTo(stream);
+ }
+ ReadWriteIOUtils.write(predictIndexList.size(), stream);
+ for (Integer index : predictIndexList) {
+ ReadWriteIOUtils.write(index, stream);
+ }
+ ReadWriteIOUtils.write(inputLength, stream);
+ ReadWriteIOUtils.write(predictLength, stream);
+ }
+
+ @Override
+ public void serialize(FileOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(TaskType.FORECAST.ordinal(), stream);
+
+ super.serialize(stream);
+ ReadWriteIOUtils.write(inputTypeList.size(), stream);
+ for (TSDataType inputType : inputTypeList) {
+ inputType.serializeTo(stream);
+ }
+ ReadWriteIOUtils.write(predictIndexList.size(), stream);
+ for (Integer index : predictIndexList) {
+ ReadWriteIOUtils.write(index, stream);
+ }
+ ReadWriteIOUtils.write(inputLength, stream);
+ ReadWriteIOUtils.write(predictLength, stream);
+ }
+}
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ModelInformation.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ModelInformation.java
index 522f609e51e..0b1eb57eebf 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ModelInformation.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/ModelInformation.java
@@ -19,7 +19,7 @@
package org.apache.iotdb.commons.model;
-import org.apache.iotdb.common.rpc.thrift.ModelTask;
+import org.apache.iotdb.common.rpc.thrift.TaskType;
import org.apache.iotdb.common.rpc.thrift.TrainingState;
import org.apache.iotdb.tsfile.utils.PublicBAOS;
import org.apache.iotdb.tsfile.utils.ReadWriteIOUtils;
@@ -32,149 +32,122 @@ import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
-import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import static org.apache.iotdb.commons.model.TrailInformation.MODEL_PATH;
+import static org.apache.iotdb.commons.model.TrialInformation.MODEL_PATH;
-public class ModelInformation {
+public abstract class ModelInformation {
private final String modelId;
- private final ModelTask modelTask;
- private final String modelType;
- private final List<String> queryExpressions;
- @Nullable private String queryFilter;
+ private final Map<String, String> options;
+ private final String datasetFetchSQL;
- private final boolean isAuto;
private TrainingState trainingState;
- @Nullable private String bestTrailId;
- private final Map<String, TrailInformation> trailMap;
+ @Nullable private String bestTrialId;
+ private final Map<String, TrialInformation> trialMap;
- public ModelInformation(
- String modelId,
- ModelTask modelTask,
- String modelType,
- boolean isAuto,
- List<String> queryExpressions,
- @Nullable String queryFilter) {
+ public static final String TASK_TYPE = "task_type";
+ public static final String MODEL_TYPE = "model_type";
+
+ protected ModelInformation(String modelId, Map<String, String> options, String datasetFetchSql) {
this.modelId = modelId;
- this.modelTask = modelTask;
- this.modelType = modelType;
- this.isAuto = isAuto;
+ this.options = options;
+ this.datasetFetchSQL = datasetFetchSql;
this.trainingState = TrainingState.PENDING;
- this.queryExpressions = queryExpressions;
- this.queryFilter = queryFilter;
- this.trailMap = new HashMap<>();
+ this.trialMap = new HashMap<>();
}
- public ModelInformation(ByteBuffer buffer) {
+ protected ModelInformation(ByteBuffer buffer) {
this.modelId = ReadWriteIOUtils.readString(buffer);
- this.modelTask = ModelTask.findByValue(ReadWriteIOUtils.readInt(buffer));
- this.modelType = ReadWriteIOUtils.readString(buffer);
-
- int listSize = ReadWriteIOUtils.readInt(buffer);
- this.queryExpressions = new ArrayList<>(listSize);
- for (int i = 0; i < listSize; i++) {
- this.queryExpressions.add(ReadWriteIOUtils.readString(buffer));
- }
- byte isNull = ReadWriteIOUtils.readByte(buffer);
- if (isNull == 1) {
- this.queryFilter = ReadWriteIOUtils.readString(buffer);
- }
+ this.options = ReadWriteIOUtils.readMap(buffer);
+ this.datasetFetchSQL = ReadWriteIOUtils.readString(buffer);
- this.isAuto = ReadWriteIOUtils.readBool(buffer);
this.trainingState = TrainingState.findByValue(ReadWriteIOUtils.readInt(buffer));
- isNull = ReadWriteIOUtils.readByte(buffer);
+ byte isNull = ReadWriteIOUtils.readByte(buffer);
if (isNull == 1) {
- this.bestTrailId = ReadWriteIOUtils.readString(buffer);
+ this.bestTrialId = ReadWriteIOUtils.readString(buffer);
}
int mapSize = ReadWriteIOUtils.readInt(buffer);
- this.trailMap = new HashMap<>();
+ this.trialMap = new HashMap<>();
for (int i = 0; i < mapSize; i++) {
- TrailInformation trailInformation = TrailInformation.deserialize(buffer);
- this.trailMap.put(trailInformation.getTrailId(), trailInformation);
+ TrialInformation trialInformation = TrialInformation.deserialize(buffer);
+ this.trialMap.put(trialInformation.getTrialId(), trialInformation);
}
}
- public ModelInformation(InputStream stream) throws IOException {
+ protected ModelInformation(InputStream stream) throws IOException {
this.modelId = ReadWriteIOUtils.readString(stream);
- this.modelTask = ModelTask.findByValue(ReadWriteIOUtils.readInt(stream));
- this.modelType = ReadWriteIOUtils.readString(stream);
-
- int listSize = ReadWriteIOUtils.readInt(stream);
- this.queryExpressions = new ArrayList<>(listSize);
- for (int i = 0; i < listSize; i++) {
- this.queryExpressions.add(ReadWriteIOUtils.readString(stream));
- }
- byte isNull = ReadWriteIOUtils.readByte(stream);
- if (isNull == 1) {
- this.queryFilter = ReadWriteIOUtils.readString(stream);
- }
+ this.options = ReadWriteIOUtils.readMap(stream);
+ this.datasetFetchSQL = ReadWriteIOUtils.readString(stream);
- this.isAuto = ReadWriteIOUtils.readBool(stream);
this.trainingState = TrainingState.findByValue(ReadWriteIOUtils.readInt(stream));
- isNull = ReadWriteIOUtils.readByte(stream);
+ byte isNull = ReadWriteIOUtils.readByte(stream);
if (isNull == 1) {
- this.bestTrailId = ReadWriteIOUtils.readString(stream);
+ this.bestTrialId = ReadWriteIOUtils.readString(stream);
}
int mapSize = ReadWriteIOUtils.readInt(stream);
- this.trailMap = new HashMap<>();
+ this.trialMap = new HashMap<>();
for (int i = 0; i < mapSize; i++) {
- TrailInformation trailInformation = TrailInformation.deserialize(stream);
- this.trailMap.put(trailInformation.getTrailId(), trailInformation);
+ TrialInformation trialInformation = TrialInformation.deserialize(stream);
+ this.trialMap.put(trialInformation.getTrialId(), trialInformation);
}
}
+ public abstract TaskType getTaskType();
+
public String getModelId() {
return modelId;
}
- public boolean isAuto() {
- return isAuto;
+ private String getModelType() {
+ return options.get(MODEL_TYPE);
}
- public List<String> getQueryExpressions() {
- return queryExpressions;
+ public Map<String, String> getOptions() {
+ return options;
}
- @Nullable
- public String getQueryFilter() {
- return queryFilter;
+ public String getDatasetFetchSql() {
+ return datasetFetchSQL;
}
- public TrailInformation getTrailInformationById(String trailId) {
- if (trailMap.containsKey(trailId)) {
- return trailMap.get(trailId);
+ public boolean available() {
+ return trainingState == TrainingState.FINISHED;
+ }
+
+ public TrialInformation getTrialInformationById(String trialId) {
+ if (trialMap.containsKey(trialId)) {
+ return trialMap.get(trialId);
}
return null;
}
- public List<TrailInformation> getAllTrailInformation() {
- return new ArrayList<>(trailMap.values());
+ public List<TrialInformation> getAllTrialInformation() {
+ return new ArrayList<>(trialMap.values());
}
public void update(String trailId, Map<String, String> modelInfo) {
- if (!trailMap.containsKey(trailId)) {
+ if (!trialMap.containsKey(trailId)) {
String modelPath = null;
if (modelInfo.containsKey(MODEL_PATH)) {
modelPath = modelInfo.get(MODEL_PATH);
modelInfo.remove(MODEL_PATH);
}
- TrailInformation trailInformation =
- new TrailInformation(trailId, new ModelHyperparameter(modelInfo), modelPath);
- trailMap.put(trailId, trailInformation);
+ TrialInformation trialInformation =
+ new TrialInformation(trailId, new ModelHyperparameter(modelInfo), modelPath);
+ trialMap.put(trailId, trialInformation);
} else {
- trailMap.get(trailId).update(modelInfo);
+ trialMap.get(trailId).update(modelInfo);
}
}
@@ -182,107 +155,107 @@ public class ModelInformation {
// TODO: add state transform validate
this.trainingState = newState;
if (bestTrailId != null) {
- this.bestTrailId = bestTrailId;
+ this.bestTrialId = bestTrailId;
+ }
+ }
+
+ public String getModelPath() {
+ if (bestTrialId != null) {
+ TrialInformation bestTrail = trialMap.get(bestTrialId);
+ return bestTrail.getModelPath();
+ } else {
+ return "UNKNOWN";
}
}
public void serialize(DataOutputStream stream) throws IOException {
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(modelTask.ordinal(), stream);
- ReadWriteIOUtils.write(modelType, stream);
- ReadWriteIOUtils.write(queryExpressions.size(), stream);
- for (String queryExpression : queryExpressions) {
- ReadWriteIOUtils.write(queryExpression, stream);
- }
- if (queryFilter == null) {
- ReadWriteIOUtils.write((byte) 0, stream);
- } else {
- ReadWriteIOUtils.write((byte) 1, stream);
- ReadWriteIOUtils.write(queryFilter, stream);
- }
+ ReadWriteIOUtils.write(options, stream);
+ ReadWriteIOUtils.write(datasetFetchSQL, stream);
- ReadWriteIOUtils.write(isAuto, stream);
ReadWriteIOUtils.write(trainingState.ordinal(), stream);
- if (bestTrailId == null) {
+ if (bestTrialId == null) {
ReadWriteIOUtils.write((byte) 0, stream);
} else {
ReadWriteIOUtils.write((byte) 1, stream);
- ReadWriteIOUtils.write(bestTrailId, stream);
+ ReadWriteIOUtils.write(bestTrialId, stream);
}
- ReadWriteIOUtils.write(trailMap.size(), stream);
- for (TrailInformation trailInformation : trailMap.values()) {
- trailInformation.serialize(stream);
+ ReadWriteIOUtils.write(trialMap.size(), stream);
+ for (TrialInformation trialInformation : trialMap.values()) {
+ trialInformation.serialize(stream);
}
}
public void serialize(FileOutputStream stream) throws IOException {
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(modelTask.ordinal(), stream);
- ReadWriteIOUtils.write(modelType, stream);
- ReadWriteIOUtils.write(queryExpressions.size(), stream);
- for (String queryExpression : queryExpressions) {
- ReadWriteIOUtils.write(queryExpression, stream);
- }
+ ReadWriteIOUtils.write(options, stream);
+ ReadWriteIOUtils.write(datasetFetchSQL, stream);
+
+ ReadWriteIOUtils.write(trainingState.ordinal(), stream);
- if (queryFilter == null) {
+ if (bestTrialId == null) {
ReadWriteIOUtils.write((byte) 0, stream);
} else {
ReadWriteIOUtils.write((byte) 1, stream);
- ReadWriteIOUtils.write(queryFilter, stream);
+ ReadWriteIOUtils.write(bestTrialId, stream);
}
- ReadWriteIOUtils.write(isAuto, stream);
- ReadWriteIOUtils.write(trainingState.ordinal(), stream);
+ ReadWriteIOUtils.write(trialMap.size(), stream);
+ for (TrialInformation trialInformation : trialMap.values()) {
+ trialInformation.serialize(stream);
+ }
+ }
- if (bestTrailId == null) {
- ReadWriteIOUtils.write((byte) 0, stream);
- } else {
- ReadWriteIOUtils.write((byte) 1, stream);
- ReadWriteIOUtils.write(bestTrailId, stream);
+ public static ModelInformation deserialize(ByteBuffer buffer) {
+ TaskType modelTask = TaskType.findByValue(ReadWriteIOUtils.readInt(buffer));
+ if (modelTask == null) {
+ throw new IllegalArgumentException();
}
- ReadWriteIOUtils.write(trailMap.size(), stream);
- for (TrailInformation trailInformation : trailMap.values()) {
- trailInformation.serialize(stream);
+ if (modelTask == TaskType.FORECAST) {
+ return new ForecastModeInformation(buffer);
}
+ throw new IllegalArgumentException("Invalid task type: " + modelTask);
}
public static ModelInformation deserialize(InputStream stream) throws IOException {
- return new ModelInformation(stream);
- }
+ TaskType modelTask = TaskType.findByValue(ReadWriteIOUtils.readInt(stream));
+ if (modelTask == null) {
+ throw new IllegalArgumentException();
+ }
- public static ModelInformation deserialize(ByteBuffer buffer) {
- return new ModelInformation(buffer);
+ if (modelTask == TaskType.FORECAST) {
+ return new ForecastModeInformation(stream);
+ }
+ throw new IllegalArgumentException("Invalid task type: " + modelTask);
}
public ByteBuffer serializeShowModelResult() throws IOException {
PublicBAOS buffer = new PublicBAOS();
DataOutputStream stream = new DataOutputStream(buffer);
ReadWriteIOUtils.write(modelId, stream);
- ReadWriteIOUtils.write(modelTask.toString(), stream);
- ReadWriteIOUtils.write(modelType, stream);
- ReadWriteIOUtils.write(Arrays.toString(queryExpressions.toArray(new String[0])), stream);
+ ReadWriteIOUtils.write(getTaskType().toString(), stream);
+ ReadWriteIOUtils.write(getModelType(), stream);
+ ReadWriteIOUtils.write(datasetFetchSQL, stream);
ReadWriteIOUtils.write(trainingState.toString(), stream);
- if (bestTrailId != null) {
- TrailInformation bestTrail = trailMap.get(bestTrailId);
- ReadWriteIOUtils.write(bestTrail.getModelPath(), stream);
-
+ if (bestTrialId != null) {
+ TrialInformation bestTrail = trialMap.get(bestTrialId);
List<String> modelHyperparameterList = bestTrail.getModelHyperparameter().toStringList();
- ReadWriteIOUtils.write(modelHyperparameterList.size(), stream);
+ ReadWriteIOUtils.write(modelHyperparameterList.size() + 1, stream);
for (String hyperparameter : modelHyperparameterList) {
ReadWriteIOUtils.write(hyperparameter, stream);
}
} else {
- ReadWriteIOUtils.write("UNKNOWN", stream);
-
- ReadWriteIOUtils.write(1, stream);
+ ReadWriteIOUtils.write(2, stream);
ReadWriteIOUtils.write("UNKNOWN", stream);
}
+ // add extra blank line to make the result more readable in cli
+ ReadWriteIOUtils.write(" ", stream);
return ByteBuffer.wrap(buffer.getBuf(), 0, buffer.size());
}
}
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrailInformation.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrialInformation.java
similarity index 78%
rename from iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrailInformation.java
rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrialInformation.java
index 56de20671e1..ee78e0a6718 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrailInformation.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/model/TrialInformation.java
@@ -30,17 +30,17 @@ import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
-public class TrailInformation {
+public class TrialInformation {
public static final String MODEL_PATH = "model_path";
- private final String trailId;
+ private final String trialId;
private final ModelHyperparameter modelHyperparameter;
private String modelPath;
- public TrailInformation(
- String trailId, ModelHyperparameter modelHyperparameter, String modelPath) {
- this.trailId = trailId;
+ public TrialInformation(
+ String trialId, ModelHyperparameter modelHyperparameter, String modelPath) {
+ this.trialId = trialId;
this.modelHyperparameter = modelHyperparameter;
this.modelPath = modelPath;
}
@@ -53,8 +53,8 @@ public class TrailInformation {
modelHyperparameter.update(modelInfo);
}
- public String getTrailId() {
- return trailId;
+ public String getTrialId() {
+ return trialId;
}
public ModelHyperparameter getModelHyperparameter() {
@@ -65,40 +65,42 @@ public class TrailInformation {
return modelPath;
}
- public ByteBuffer serializeShowTrailResult() throws IOException {
+ public ByteBuffer serializeShowTrialResult() throws IOException {
PublicBAOS buffer = new PublicBAOS();
DataOutputStream stream = new DataOutputStream(buffer);
- ReadWriteIOUtils.write(trailId, stream);
+ ReadWriteIOUtils.write(trialId, stream);
ReadWriteIOUtils.write(modelPath, stream);
List<String> modelHyperparameterList = modelHyperparameter.toStringList();
- ReadWriteIOUtils.write(modelHyperparameterList.size(), stream);
+ ReadWriteIOUtils.write(modelHyperparameterList.size() + 1, stream);
for (String hyperparameter : modelHyperparameterList) {
ReadWriteIOUtils.write(hyperparameter, stream);
}
+ // add a line break to make the result more readable.
+ ReadWriteIOUtils.write(" ", stream);
return ByteBuffer.wrap(buffer.getBuf(), 0, buffer.size());
}
public void serialize(DataOutputStream stream) throws IOException {
- ReadWriteIOUtils.write(trailId, stream);
+ ReadWriteIOUtils.write(trialId, stream);
modelHyperparameter.serialize(stream);
ReadWriteIOUtils.write(modelPath, stream);
}
public void serialize(FileOutputStream stream) throws IOException {
- ReadWriteIOUtils.write(trailId, stream);
+ ReadWriteIOUtils.write(trialId, stream);
modelHyperparameter.serialize(stream);
ReadWriteIOUtils.write(modelPath, stream);
}
- public static TrailInformation deserialize(ByteBuffer buffer) {
- return new TrailInformation(
+ public static TrialInformation deserialize(ByteBuffer buffer) {
+ return new TrialInformation(
ReadWriteIOUtils.readString(buffer),
ModelHyperparameter.deserialize(buffer),
ReadWriteIOUtils.readString(buffer));
}
- public static TrailInformation deserialize(InputStream stream) throws IOException {
- return new TrailInformation(
+ public static TrialInformation deserialize(InputStream stream) throws IOException {
+ return new TrialInformation(
ReadWriteIOUtils.readString(stream),
ModelHyperparameter.deserialize(stream),
ReadWriteIOUtils.readString(stream));
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/builtin/ModelInferenceFunction.java
similarity index 52%
copy from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
copy to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/builtin/ModelInferenceFunction.java
index 336bbcce0bb..f52a0348af6 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionType.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/builtin/ModelInferenceFunction.java
@@ -17,11 +17,33 @@
* under the License.
*/
-package org.apache.iotdb.db.queryengine.plan.expression.multi;
+package org.apache.iotdb.commons.udf.builtin;
-/** */
-public enum FunctionType {
- AGGREGATION_FUNCTION,
- BUILT_IN_SCALAR_FUNCTION,
- UDF
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public enum ModelInferenceFunction {
+ FORECAST("forecast");
+
+ private final String functionName;
+
+ ModelInferenceFunction(String functionName) {
+ this.functionName = functionName;
+ }
+
+ public String getFunctionName() {
+ return functionName;
+ }
+
+ private static final Set<String> NATIVE_FUNCTION_NAMES =
+ new HashSet<>(
+ Arrays.stream(ModelInferenceFunction.values())
+ .map(ModelInferenceFunction::getFunctionName)
+ .collect(Collectors.toList()));
+
+ public static Set<String> getNativeFunctionNames() {
+ return NATIVE_FUNCTION_NAMES;
+ }
}
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFManagementService.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFManagementService.java
index fbdf8684e91..e7ecb32f008 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFManagementService.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFManagementService.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.commons.udf.service;
import org.apache.iotdb.commons.udf.UDFInformation;
import org.apache.iotdb.commons.udf.UDFTable;
import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
+import org.apache.iotdb.commons.udf.builtin.ModelInferenceFunction;
import org.apache.iotdb.commons.utils.TestOnly;
import org.apache.iotdb.udf.api.UDF;
import org.apache.iotdb.udf.api.UDTF;
@@ -111,8 +112,26 @@ public class UDFManagementService {
throw new UDFManagementException(errorMessage);
}
+ private void checkIsModelInferenceFunctionName(UDFInformation udfInformation)
+ throws UDFManagementException {
+ String functionName = udfInformation.getFunctionName();
+ String className = udfInformation.getClassName();
+ if (!ModelInferenceFunction.getNativeFunctionNames().contains(functionName.toLowerCase())) {
+ return;
+ }
+
+ String errorMessage =
+ String.format(
+ "Failed to register UDF %s(%s), because the given function name conflicts with the ML model inference function name",
+ functionName, className);
+
+ LOGGER.warn(errorMessage);
+ throw new UDFManagementException(errorMessage);
+ }
+
private void checkIfRegistered(UDFInformation udfInformation) throws UDFManagementException {
checkIsBuiltInAggregationFunctionName(udfInformation);
+ checkIsModelInferenceFunctionName(udfInformation);
String functionName = udfInformation.getFunctionName();
String className = udfInformation.getClassName();
UDFInformation information = udfTable.getUDFInformation(functionName);
diff --git a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/file/metadata/enums/TSDataType.java b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/file/metadata/enums/TSDataType.java
index 102b52ef0b3..8c5ffc7e368 100644
--- a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/file/metadata/enums/TSDataType.java
+++ b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/file/metadata/enums/TSDataType.java
@@ -22,7 +22,9 @@ package org.apache.iotdb.tsfile.file.metadata.enums;
import org.apache.iotdb.tsfile.exception.write.UnSupportedDataTypeException;
import java.io.DataOutputStream;
+import java.io.FileOutputStream;
import java.io.IOException;
+import java.io.InputStream;
import java.nio.ByteBuffer;
public enum TSDataType {
@@ -97,6 +99,10 @@ public enum TSDataType {
return deserialize(buffer.get());
}
+ public static TSDataType deserializeFrom(InputStream stream) throws IOException {
+ return deserialize((byte) stream.read());
+ }
+
public static int getSerializedSize() {
return Byte.BYTES;
}
@@ -109,6 +115,10 @@ public enum TSDataType {
outputStream.write(serialize());
}
+ public void serializeTo(FileOutputStream outputStream) throws IOException {
+ outputStream.write(serialize());
+ }
+
public int getDataTypeSize() {
switch (this) {
case BOOLEAN:
diff --git a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlock.java b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlock.java
index a22793fa486..677e9aea89d 100644
--- a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlock.java
+++ b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlock.java
@@ -137,14 +137,16 @@ public class TsBlock {
length, (TimeColumn) timeColumn.getRegion(positionOffset, length), slicedColumns);
}
- public TsBlock appendValueColumn(Column column) {
- requireNonNull(column, "Column is null");
- if (positionCount != column.getPositionCount()) {
- throw new IllegalArgumentException("Block does not have same position count");
+ public TsBlock appendValueColumns(Column[] columns) {
+ Column[] newBlocks = Arrays.copyOf(valueColumns, valueColumns.length + columns.length);
+ int newColumnIndex = valueColumns.length;
+ for (Column column : columns) {
+ requireNonNull(column, "Column is null");
+ if (positionCount != column.getPositionCount()) {
+ throw new IllegalArgumentException("Block does not have same position count");
+ }
+ newBlocks[newColumnIndex++] = column;
}
-
- Column[] newBlocks = Arrays.copyOf(valueColumns, valueColumns.length + 1);
- newBlocks[valueColumns.length] = column;
return wrapBlocksWithoutCopy(positionCount, timeColumn, newBlocks);
}
@@ -469,8 +471,8 @@ public class TsBlock {
}
private boolean isCurrentValueAllNull() {
- for (int i = 0; i < valueColumns.length; i++) {
- if (!valueColumns[i].isNull(rowIndex)) {
+ for (Column valueColumn : valueColumns) {
+ if (!valueColumn.isNull(rowIndex)) {
return false;
}
}
diff --git a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlockBuilder.java b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlockBuilder.java
index 29370d2735f..883e5a2133a 100644
--- a/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlockBuilder.java
+++ b/iotdb-core/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/block/TsBlockBuilder.java
@@ -82,6 +82,7 @@ public class TsBlockBuilder {
res.timeColumnBuilder =
new TimeColumnBuilder(
res.tsBlockBuilderStatus.createColumnBuilderStatus(), DEFAULT_INITIAL_EXPECTED_ENTRIES);
+ res.valueColumnBuilders = new ColumnBuilder[0];
return res;
}
@@ -286,9 +287,6 @@ public class TsBlockBuilder {
}
public TsBlock build() {
- if (valueColumnBuilders.length == 0) {
- return new TsBlock(declaredPositions);
- }
TimeColumn timeColumn = (TimeColumn) timeColumnBuilder.build();
if (timeColumn.getPositionCount() != declaredPositions) {
throw new IllegalStateException(
diff --git a/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift b/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift
index d602c505319..75e795bf1cf 100644
--- a/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift
+++ b/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift
@@ -191,6 +191,7 @@ enum TrainingState {
DROPPING
}
-enum ModelTask {
- FORECAST
+enum TaskType {
+ FORECAST,
+ ANOMALY_DETECTION,
}
\ No newline at end of file
diff --git a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
index a19f5536f5f..382a2783991 100644
--- a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
+++ b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
@@ -728,12 +728,10 @@ struct TUnsetSchemaTemplateReq{
struct TCreateModelReq {
1: required string modelId
- 2: required common.ModelTask modelTask
- 3: required string modelType
- 4: required list<string> queryExpressions
- 5: optional string queryFilter
- 6: required bool isAuto
- 7: required map<string, string> modelConfigs
+ 2: required common.TaskType taskType
+ 3: required map<string, string> options
+ 4: required map<string, string> hyperparameters
+ 5: required string datasetFetchSQL
}
struct TDropModelReq {
@@ -749,26 +747,35 @@ struct TShowModelResp {
2: required list<binary> modelInfoList
}
-struct TShowTrailReq {
+struct TShowTrialReq {
1: required string modelId
- 2: optional string trailId
+ 2: optional string trialId
}
-struct TShowTrailResp {
+struct TGetModelInfoReq {
+ 1: required string modelId
+}
+
+struct TShowTrialResp {
1: required common.TSStatus status
- 2: required list<binary> trailInfoList
+ 2: required list<binary> trialInfoList
}
struct TUpdateModelInfoReq {
1: required string modelId
- 2: required string trailId
+ 2: required string trialId
3: required map<string, string> modelInfo
}
struct TUpdateModelStateReq {
1: required string modelId
2: required common.TrainingState state
- 3: optional string bestTrailId
+ 3: optional string bestTrialId
+}
+
+struct TGetModelInfoResp {
+ 1: required common.TSStatus status
+ 2: optional binary modelInfo
}
// ====================================================
@@ -1371,9 +1378,9 @@ service IConfigNodeRPCService {
TShowModelResp showModel(TShowModelReq req)
/**
- * Return the trail table
+ * Return the trial table
*/
- TShowTrailResp showTrail(TShowTrailReq req)
+ TShowTrialResp showTrial(TShowTrialReq req)
/**
* Update the model info
@@ -1389,6 +1396,11 @@ service IConfigNodeRPCService {
*/
common.TSStatus updateModelState(TUpdateModelStateReq req)
+ /**
+ * Return the model info by model_id
+ */
+ TGetModelInfoResp getModelInfo(TGetModelInfoReq req)
+
// ======================================================
// Quota
// ======================================================
diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
index 1268c399c5f..4c4fc95d932 100644
--- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
+++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
@@ -465,10 +465,9 @@ struct TFetchMoreDataResp{
}
struct TFetchTimeseriesReq {
- 1: required list<string> queryExpressions
- 2: optional string queryFilter
- 3: optional i32 fetchSize
- 4: optional i64 timeout
+ 1: required string queryBody
+ 2: optional i32 fetchSize
+ 3: optional i64 timeout
}
struct TFetchTimeseriesResp {
diff --git a/iotdb-protocol/thrift-mlnode/src/main/thrift/mlnode.thrift b/iotdb-protocol/thrift-mlnode/src/main/thrift/mlnode.thrift
index abadc795768..3d7f2fe0a71 100644
--- a/iotdb-protocol/thrift-mlnode/src/main/thrift/mlnode.thrift
+++ b/iotdb-protocol/thrift-mlnode/src/main/thrift/mlnode.thrift
@@ -23,10 +23,9 @@ namespace py iotdb.thrift.mlnode
struct TCreateTrainingTaskReq {
1: required string modelId
- 2: required bool isAuto
- 3: required map<string, string> modelConfigs
- 4: required list<string> queryExpressions
- 5: optional string queryFilter
+ 2: required map<string, string> options
+ 3: required map<string, string> hyperparameters
+ 4: required string datasetFetchSQL
}
struct TDeleteModelReq {
@@ -36,7 +35,10 @@ struct TDeleteModelReq {
struct TForecastReq {
1: required string modelPath
- 2: required binary dataset
+ 2: required binary inputData
+ 3: required list<string> inputTypeList
+ 4: required list<string> inputColumnNameList
+ 5: required i32 predictLength
}
struct TForecastResp {