org.apache.iotdb
pipe-api
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java
index 1b291eb7c335..9f6355b99ebf 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java
@@ -432,6 +432,9 @@ public class IoTDBConfig {
/** Compact the unsequence files into the overlapped sequence files */
private volatile boolean enableCrossSpaceCompaction = true;
+ /** Enable the service for AINode */
+ private boolean enableAINodeService = false;
+
/** The buffer for sort operation */
private long sortBufferSize = 1024 * 1024L;
@@ -910,6 +913,9 @@ public class IoTDBConfig {
/** Internal port for coordinator */
private int internalPort = 10730;
+ /** Port for AINode */
+ private int aiNodePort = 10780;
+
/** Internal port for dataRegion consensus protocol */
private int dataRegionConsensusPort = 10760;
@@ -2848,6 +2854,14 @@ public void setEnableCrossSpaceCompaction(boolean enableCrossSpaceCompaction) {
this.enableCrossSpaceCompaction = enableCrossSpaceCompaction;
}
+ public boolean isEnableAINodeService() {
+ return enableAINodeService;
+ }
+
+ public void setEnableAINodeService(boolean enableAINodeService) {
+ this.enableAINodeService = enableAINodeService;
+ }
+
public InnerSequenceCompactionSelector getInnerSequenceCompactionSelector() {
return innerSequenceCompactionSelector;
}
@@ -3126,6 +3140,14 @@ public void setInternalPort(int internalPort) {
this.internalPort = internalPort;
}
+ public int getAINodePort() {
+ return aiNodePort;
+ }
+
+ public void setAINodePort(int aiNodePort) {
+ this.aiNodePort = aiNodePort;
+ }
+
public int getDataRegionConsensusPort() {
return dataRegionConsensusPort;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
index d8da4478ac30..5eaa88e3ee5d 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
@@ -287,6 +287,19 @@ public void loadProperties(Properties properties) throws BadNodeUrlException, IO
.getProperty(IoTDBConstant.DN_RPC_PORT, Integer.toString(conf.getRpcPort()))
.trim()));
+ conf.setEnableAINodeService(
+ Boolean.parseBoolean(
+ properties
+ .getProperty(
+ "enable_ainode_rpc_service", Boolean.toString(conf.isEnableAINodeService()))
+ .trim()));
+
+ conf.setAINodePort(
+ Integer.parseInt(
+ properties
+ .getProperty("ainode_rpc_port", Integer.toString(conf.getAINodePort()))
+ .trim()));
+
conf.setBufferedArraysMemoryProportion(
Double.parseDouble(
properties
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java
new file mode 100644
index 000000000000..03402d30c643
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java
@@ -0,0 +1,28 @@
+/*
+ * 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.exception.ainode;
+
+import org.apache.iotdb.rpc.TSStatusCode;
+
+public class GetModelInfoException extends ModelException {
+ public GetModelInfoException(String message) {
+ super(message, TSStatusCode.GET_MODEL_INFO_ERROR);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelException.java
new file mode 100644
index 000000000000..4a007e7048ce
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelException.java
@@ -0,0 +1,35 @@
+/*
+ * 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.exception.ainode;
+
+import org.apache.iotdb.rpc.TSStatusCode;
+
+public class ModelException extends RuntimeException {
+ TSStatusCode statusCode;
+
+ public ModelException(String message, TSStatusCode code) {
+ super(message);
+ this.statusCode = code;
+ }
+
+ public TSStatusCode getStatusCode() {
+ return statusCode;
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java
new file mode 100644
index 000000000000..38a5105cded1
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java
@@ -0,0 +1,28 @@
+/*
+ * 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.exception.ainode;
+
+import org.apache.iotdb.rpc.TSStatusCode;
+
+public class ModelNotFoundException extends ModelException {
+ public ModelNotFoundException(String message) {
+ super(message, TSStatusCode.MODEL_NOT_FOUND_ERROR);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/runtime/ModelInferenceProcessException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/runtime/ModelInferenceProcessException.java
new file mode 100644
index 000000000000..586c624a8d33
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/runtime/ModelInferenceProcessException.java
@@ -0,0 +1,27 @@
+/*
+ * 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.exception.runtime;
+
+public class ModelInferenceProcessException extends RuntimeException {
+
+ public ModelInferenceProcessException(String message) {
+ super(message);
+ }
+}
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 122ad5d7c37e..d146712ecf1b 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
@@ -40,6 +40,12 @@
import org.apache.iotdb.commons.client.sync.SyncThriftClientWithErrorHandler;
import org.apache.iotdb.commons.consensus.ConfigRegionId;
import org.apache.iotdb.confignode.rpc.thrift.IConfigNodeRPCService;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeConfigurationResp;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeRegisterReq;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeRegisterResp;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeRemoveReq;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeRestartReq;
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeRestartResp;
import org.apache.iotdb.confignode.rpc.thrift.TAddConsensusGroupReq;
import org.apache.iotdb.confignode.rpc.thrift.TAlterLogicalViewReq;
import org.apache.iotdb.confignode.rpc.thrift.TAlterPipeReq;
@@ -59,6 +65,7 @@
import org.apache.iotdb.confignode.rpc.thrift.TCreateCQReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreateConsumerReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreateFunctionReq;
+import org.apache.iotdb.confignode.rpc.thrift.TCreateModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreatePipePluginReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreateSchemaTemplateReq;
@@ -82,6 +89,7 @@
import org.apache.iotdb.confignode.rpc.thrift.TDeleteTimeSeriesReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropCQReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropFunctionReq;
+import org.apache.iotdb.confignode.rpc.thrift.TDropModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropPipePluginReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropTopicReq;
@@ -96,6 +104,8 @@
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.TGetPathsSetTemplatesReq;
import org.apache.iotdb.confignode.rpc.thrift.TGetPathsSetTemplatesResp;
import org.apache.iotdb.confignode.rpc.thrift.TGetPipePluginTableResp;
@@ -123,11 +133,14 @@
import org.apache.iotdb.confignode.rpc.thrift.TSetSchemaReplicationFactorReq;
import org.apache.iotdb.confignode.rpc.thrift.TSetSchemaTemplateReq;
import org.apache.iotdb.confignode.rpc.thrift.TSetTimePartitionIntervalReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowAINodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowCQResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowConfigNodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowDatabaseResp;
+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.TShowRegionReq;
@@ -447,6 +460,33 @@ public TDataNodeRestartResp restartDataNode(TDataNodeRestartReq req) throws TExc
() -> client.restartDataNode(req), resp -> !updateConfigNodeLeader(resp.status));
}
+ @Override
+ public TAINodeRegisterResp registerAINode(TAINodeRegisterReq req) throws TException {
+ throw new UnsupportedOperationException("RegisterAINode method is not supported in datanode");
+ }
+
+ @Override
+ public TAINodeRestartResp restartAINode(TAINodeRestartReq req) throws TException {
+ throw new UnsupportedOperationException("RestartAINode method is not supported in datanode");
+ }
+
+ @Override
+ public TSStatus removeAINode(TAINodeRemoveReq req) throws TException {
+ throw new UnsupportedOperationException("RemoveAINode method is not supported in datanode");
+ }
+
+ @Override
+ public TShowAINodesResp showAINodes() throws TException {
+ return executeRemoteCallWithRetry(
+ () -> client.showAINodes(), resp -> !updateConfigNodeLeader(resp.status));
+ }
+
+ @Override
+ public TAINodeConfigurationResp getAINodeConfiguration(int aiNodeId) throws TException {
+ throw new UnsupportedOperationException(
+ "GetAINodeConfiguration method is not supported in datanode");
+ }
+
@Override
public TDataNodeRemoveResp removeDataNode(TDataNodeRemoveReq req) throws TException {
return executeRemoteCallWithRetry(
@@ -1139,6 +1179,30 @@ public TShowCQResp showCQ() throws TException {
() -> client.showCQ(), resp -> !updateConfigNodeLeader(resp.status));
}
+ @Override
+ public TSStatus createModel(TCreateModelReq req) throws TException {
+ return executeRemoteCallWithRetry(
+ () -> client.createModel(req), status -> !updateConfigNodeLeader(status));
+ }
+
+ @Override
+ public TSStatus dropModel(TDropModelReq req) throws TException {
+ return executeRemoteCallWithRetry(
+ () -> client.dropModel(req), status -> !updateConfigNodeLeader(status));
+ }
+
+ @Override
+ public TShowModelResp showModel(TShowModelReq req) throws TException {
+ return executeRemoteCallWithRetry(
+ () -> client.showModel(req), resp -> !updateConfigNodeLeader(resp.status));
+ }
+
+ @Override
+ public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) throws TException {
+ return executeRemoteCallWithRetry(
+ () -> client.getModelInfo(req), resp -> !updateConfigNodeLeader(resp.getStatus()));
+ }
+
@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/handler/AINodeRPCServiceThriftHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/AINodeRPCServiceThriftHandler.java
new file mode 100644
index 000000000000..c5969f8678a3
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/AINodeRPCServiceThriftHandler.java
@@ -0,0 +1,60 @@
+/*
+ * 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.protocol.thrift.handler;
+
+import org.apache.iotdb.db.protocol.thrift.impl.IAINodeRPCServiceWithHandler;
+
+import org.apache.thrift.protocol.TProtocol;
+import org.apache.thrift.server.ServerContext;
+import org.apache.thrift.server.TServerEventHandler;
+import org.apache.thrift.transport.TTransport;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+public class AINodeRPCServiceThriftHandler implements TServerEventHandler {
+
+ private final AtomicLong thriftConnectionNumber = new AtomicLong(0);
+ private final IAINodeRPCServiceWithHandler eventHandler;
+
+ public AINodeRPCServiceThriftHandler(IAINodeRPCServiceWithHandler eventHandler) {
+ this.eventHandler = eventHandler;
+ }
+
+ @Override
+ public ServerContext createContext(TProtocol in, TProtocol out) {
+ thriftConnectionNumber.incrementAndGet();
+ return null;
+ }
+
+ @Override
+ public void deleteContext(ServerContext arg0, TProtocol in, TProtocol out) {
+ thriftConnectionNumber.decrementAndGet();
+ eventHandler.handleExit();
+ }
+
+ @Override
+ public void preServe() {
+ // do nothing
+ }
+
+ @Override
+ public void processContext(
+ ServerContext serverContext, TTransport tTransport, TTransport tTransport1) {
+ // do nothing
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/AINodeRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/AINodeRPCServiceImpl.java
new file mode 100644
index 000000000000..68e492cdf462
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/AINodeRPCServiceImpl.java
@@ -0,0 +1,187 @@
+/*
+ * 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.protocol.thrift.impl;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.conf.IoTDBConstant.ClientVersion;
+import org.apache.iotdb.db.protocol.session.IClientSession;
+import org.apache.iotdb.db.protocol.session.InternalClientSession;
+import org.apache.iotdb.db.protocol.session.SessionManager;
+import org.apache.iotdb.db.protocol.thrift.OperationType;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
+import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
+import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
+import org.apache.iotdb.db.queryengine.plan.analyze.schema.ClusterSchemaFetcher;
+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.utils.ErrorHandlingUtils;
+import org.apache.iotdb.db.utils.QueryDataSetUtils;
+import org.apache.iotdb.db.utils.SetThreadName;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchMoreDataReq;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchMoreDataResp;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchTimeseriesReq;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchTimeseriesResp;
+import org.apache.iotdb.rpc.RpcUtils;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.apache.thrift.TException;
+import org.apache.tsfile.utils.Pair;
+
+import java.nio.ByteBuffer;
+import java.time.ZoneId;
+import java.util.List;
+
+public class AINodeRPCServiceImpl implements IAINodeRPCServiceWithHandler {
+
+ public static final String AI_METRICS_PATH_PREFIX = "root.__system.AI.exp";
+
+ private static final SessionManager SESSION_MANAGER = SessionManager.getInstance();
+
+ private static final Coordinator COORDINATOR = Coordinator.getInstance();
+
+ private final IPartitionFetcher partitionFetcher;
+
+ private final ISchemaFetcher schemaFetcher;
+
+ private final IClientSession session;
+
+ public AINodeRPCServiceImpl() {
+ super();
+ partitionFetcher = ClusterPartitionFetcher.getInstance();
+ schemaFetcher = ClusterSchemaFetcher.getInstance();
+ session = new InternalClientSession("AINodeService");
+ SESSION_MANAGER.registerSession(session);
+ SESSION_MANAGER.supplySession(session, "AINode", ZoneId.systemDefault(), ClientVersion.V_1_0);
+ }
+
+ @Override
+ public TFetchTimeseriesResp fetchTimeseries(TFetchTimeseriesReq req) throws TException {
+ boolean finished = false;
+ TFetchTimeseriesResp resp = new TFetchTimeseriesResp();
+ Throwable t = null;
+ try {
+
+ 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));
+ ExecutionResult result =
+ COORDINATOR.executeForTreeModel(
+ s,
+ queryId,
+ SESSION_MANAGER.getSessionInfo(session),
+ "",
+ partitionFetcher,
+ schemaFetcher,
+ req.getTimeout());
+
+ if (result.status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()
+ && result.status.code != TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()) {
+ resp.setStatus(result.status);
+ return resp;
+ }
+
+ IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId);
+
+ try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) {
+
+ DatasetHeader header = queryExecution.getDatasetHeader();
+ resp.setStatus(result.status);
+ resp.setColumnNameList(header.getRespColumns());
+ resp.setColumnTypeList(header.getRespDataTypeList());
+ resp.setColumnNameIndexMap(header.getColumnNameIndexMap());
+ resp.setQueryId(queryId);
+
+ Pair, Boolean> pair =
+ QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, req.fetchSize);
+ resp.setTsDataset(pair.left);
+ finished = pair.right;
+ resp.setHasMoreData(!finished);
+ return resp;
+ }
+ } catch (Exception e) {
+ finished = true;
+ t = e;
+ resp.setStatus(ErrorHandlingUtils.onQueryException(e, OperationType.EXECUTE_STATEMENT));
+ return resp;
+ } catch (Error error) {
+ t = error;
+ throw error;
+ } finally {
+ if (finished) {
+ COORDINATOR.cleanupQueryExecution(resp.queryId, req, t);
+ }
+ }
+ }
+
+ @Override
+ public TFetchMoreDataResp fetchMoreData(TFetchMoreDataReq req) throws TException {
+ TFetchMoreDataResp resp = new TFetchMoreDataResp();
+ boolean finished = false;
+ Throwable t = null;
+ try {
+ IQueryExecution queryExecution = COORDINATOR.getQueryExecution(req.queryId);
+ resp.setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+
+ if (queryExecution == null) {
+ resp.setHasMoreData(false);
+ return resp;
+ }
+
+ try (SetThreadName queryName = new SetThreadName(queryExecution.getQueryId())) {
+ Pair, Boolean> pair =
+ QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, req.fetchSize);
+ List result = pair.left;
+ finished = pair.right;
+ resp.setTsDataset(result);
+ resp.setHasMoreData(!finished);
+ return resp;
+ }
+ } catch (Exception e) {
+ finished = true;
+ t = e;
+ resp.setStatus(ErrorHandlingUtils.onQueryException(e, OperationType.FETCH_RESULTS));
+ return resp;
+ } catch (Error error) {
+ t = error;
+ throw error;
+ } finally {
+ if (finished) {
+ COORDINATOR.cleanupQueryExecution(req.queryId, req, t);
+ }
+ }
+ }
+
+ @Override
+ public void handleExit() {
+ SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
index 7755bb99e31c..3604f0c00ddf 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
@@ -24,6 +24,7 @@
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TFlushReq;
+import org.apache.iotdb.common.rpc.thrift.TLoadSample;
import org.apache.iotdb.common.rpc.thrift.TNodeLocations;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.common.rpc.thrift.TSender;
@@ -204,7 +205,6 @@
import org.apache.iotdb.mpp.rpc.thrift.TInvalidatePermissionCacheReq;
import org.apache.iotdb.mpp.rpc.thrift.TLoadCommandReq;
import org.apache.iotdb.mpp.rpc.thrift.TLoadResp;
-import org.apache.iotdb.mpp.rpc.thrift.TLoadSample;
import org.apache.iotdb.mpp.rpc.thrift.TMaintainPeerReq;
import org.apache.iotdb.mpp.rpc.thrift.TPipeHeartbeatReq;
import org.apache.iotdb.mpp.rpc.thrift.TPipeHeartbeatResp;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/IAINodeRPCServiceWithHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/IAINodeRPCServiceWithHandler.java
new file mode 100644
index 000000000000..7d9df50ac3ce
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/IAINodeRPCServiceWithHandler.java
@@ -0,0 +1,26 @@
+/*
+ * 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.protocol.thrift.impl;
+
+import org.apache.iotdb.mpp.rpc.thrift.IAINodeInternalRPCService;
+
+public interface IAINodeRPCServiceWithHandler extends IAINodeInternalRPCService.Iface {
+ void handleExit();
+}
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 f4fd39bf504d..088dcdabde80 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
@@ -46,6 +46,7 @@ private ColumnHeaderConstant() {
public static final String COMPRESSION = "Compression";
public static final String TAGS = "Tags";
public static final String ATTRIBUTES = "Attributes";
+ public static final String NOTES = "Notes";
public static final String DEADBAND = "Deadband";
public static final String DEADBAND_PARAMETERS = "DeadbandParameters";
public static final String IS_ALIGNED = "IsAligned";
@@ -109,6 +110,8 @@ private ColumnHeaderConstant() {
public static final String TRIGGER_NAME = "TriggerName";
public static final String EVENT = "Event";
public static final String STATE = "State";
+ public static final String MODEL_TYPE = "ModelType";
+ public static final String CONFIGS = "Configs";
public static final String PATH_PATTERN = "PathPattern";
public static final String CLASS_NAME = "ClassName";
@@ -128,6 +131,7 @@ private ColumnHeaderConstant() {
// show cluster status
public static final String NODE_TYPE_CONFIG_NODE = "ConfigNode";
public static final String NODE_TYPE_DATA_NODE = "DataNode";
+ public static final String NODE_TYPE_AI_NODE = "AINode";
public static final String COLUMN_CLUSTER_NAME = "ClusterName";
public static final String CONFIG_NODE_CONSENSUS_PROTOCOL_CLASS =
"ConfigNodeConsensusProtocolClass";
@@ -200,6 +204,9 @@ private ColumnHeaderConstant() {
public static final String USER = "User";
public static final String READ_WRITE = "Read/Write";
+ // column names for show models/trials
+ public static final String MODEL_ID = "ModelId";
+
// column names for views (e.g. logical view)
public static final String VIEW_TYPE = "ViewType";
public static final String SOURCE = "Source";
@@ -329,6 +336,13 @@ private ColumnHeaderConstant() {
new ColumnHeader(ROLE, TSDataType.TEXT),
new ColumnHeader(CREATE_TIME, TSDataType.TEXT));
+ public static final List showAINodesColumnHeaders =
+ ImmutableList.of(
+ new ColumnHeader(NODE_ID, TSDataType.INT32),
+ new ColumnHeader(STATUS, TSDataType.TEXT),
+ new ColumnHeader(RPC_ADDRESS, TSDataType.TEXT),
+ new ColumnHeader(RPC_PORT, TSDataType.INT32));
+
public static final List showDataNodesColumnHeaders =
ImmutableList.of(
new ColumnHeader(NODE_ID, TSDataType.INT32),
@@ -497,6 +511,14 @@ private ColumnHeaderConstant() {
new ColumnHeader(LIMIT, TSDataType.TEXT),
new ColumnHeader(READ_WRITE, TSDataType.TEXT));
+ public static final List showModelsColumnHeaders =
+ ImmutableList.of(
+ new ColumnHeader(MODEL_ID, TSDataType.TEXT),
+ new ColumnHeader(MODEL_TYPE, TSDataType.TEXT),
+ new ColumnHeader(STATE, TSDataType.TEXT),
+ new ColumnHeader(CONFIGS, TSDataType.TEXT),
+ new ColumnHeader(NOTES, TSDataType.TEXT));
+
public static final List showLogicalViewColumnHeaders =
ImmutableList.of(
new ColumnHeader(TIMESERIES, 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 9c0f89b19a57..132dafd246d9 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
@@ -119,6 +119,10 @@ public static DatasetHeader getShowRegionHeader() {
return new DatasetHeader(ColumnHeaderConstant.showRegionColumnHeaders, true);
}
+ public static DatasetHeader getShowAINodesHeader() {
+ return new DatasetHeader(ColumnHeaderConstant.showAINodesColumnHeaders, true);
+ }
+
public static DatasetHeader getShowDataNodesHeader() {
return new DatasetHeader(ColumnHeaderConstant.showDataNodesColumnHeaders, true);
}
@@ -201,6 +205,10 @@ public static DatasetHeader getShowThrottleQuotaHeader() {
return new DatasetHeader(ColumnHeaderConstant.showThrottleQuotaColumnHeaders, true);
}
+ public static DatasetHeader getShowModelsHeader() {
+ return new DatasetHeader(ColumnHeaderConstant.showModelsColumnHeaders, true);
+ }
+
public static DatasetHeader getShowLogicalViewHeader() {
return new DatasetHeader(ColumnHeaderConstant.showLogicalViewColumnHeaders, true);
}
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 40a120da7250..a82e5a6d92eb 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
@@ -23,6 +23,7 @@
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.commons.exception.IoTDBException;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
@@ -82,6 +83,7 @@ public class FragmentInstanceManager {
private final Duration infoCacheTime;
private final ExecutorService intoOperationExecutor;
+ private final ExecutorService modelInferenceExecutor;
private final MPPDataExchangeManager exchangeManager =
MPPDataExchangeService.getInstance().getMPPDataExchangeManager();
@@ -119,6 +121,11 @@ private FragmentInstanceManager() {
IoTDBThreadPoolFactory.newFixedThreadPool(
IoTDBDescriptor.getInstance().getConfig().getIntoOperationExecutionThreadCount(),
"into-operation-executor");
+
+ this.modelInferenceExecutor =
+ IoTDBThreadPoolFactory.newFixedThreadPool(
+ CommonDescriptor.getInstance().getConfig().getModelInferenceExecutionThreadCount(),
+ "model-inference-executor");
}
@SuppressWarnings("squid:S1181")
@@ -424,6 +431,10 @@ public ExecutorService getIntoOperationExecutor() {
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/process/AI/InferenceOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AI/InferenceOperator.java
new file mode 100644
index 000000000000..ab8cf811af7e
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AI/InferenceOperator.java
@@ -0,0 +1,337 @@
+/*
+ * 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.AI;
+
+import org.apache.iotdb.ainode.rpc.thrift.TInferenceResp;
+import org.apache.iotdb.ainode.rpc.thrift.TWindowParams;
+import org.apache.iotdb.commons.client.ainode.AINodeClient;
+import org.apache.iotdb.commons.client.ainode.AINodeClientManager;
+import org.apache.iotdb.db.exception.runtime.ModelInferenceProcessException;
+import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
+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.queryengine.execution.operator.window.ainode.BottomInferenceWindowParameter;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.CountInferenceWindowParameter;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindowType;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.model.ModelInferenceDescriptor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.read.common.block.column.TimeColumnBuilder;
+import org.apache.tsfile.read.common.block.column.TsBlockSerde;
+import org.apache.tsfile.utils.RamUsageEstimator;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.stream.Collectors;
+
+import static com.google.common.util.concurrent.Futures.successfulAsList;
+
+public class InferenceOperator implements ProcessOperator {
+
+ private static final long INSTANCE_SIZE =
+ RamUsageEstimator.shallowSizeOfInstance(InferenceOperator.class);
+
+ private final OperatorContext operatorContext;
+ private final Operator child;
+ private final ModelInferenceDescriptor modelInferenceDescriptor;
+
+ private final TsBlockBuilder inputTsBlockBuilder;
+
+ private final ExecutorService modelInferenceExecutor;
+ private ListenableFuture inferenceExecutionFuture;
+
+ private boolean finished = false;
+
+ private final long maxRetainedSize;
+ private final long maxReturnSize;
+ private final List inputColumnNames;
+ private final List targetColumnNames;
+ private long totalRow;
+ private int resultIndex = 0;
+ private List results;
+ private final TsBlockSerde serde = new TsBlockSerde();
+ private InferenceWindowType windowType = null;
+
+ public InferenceOperator(
+ OperatorContext operatorContext,
+ Operator child,
+ ModelInferenceDescriptor modelInferenceDescriptor,
+ ExecutorService modelInferenceExecutor,
+ List targetColumnNames,
+ List inputColumnNames,
+ long maxRetainedSize,
+ long maxReturnSize) {
+ this.operatorContext = operatorContext;
+ this.child = child;
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ this.inputTsBlockBuilder =
+ new TsBlockBuilder(
+ Arrays.asList(modelInferenceDescriptor.getModelInformation().getInputDataType()));
+ this.modelInferenceExecutor = modelInferenceExecutor;
+ this.targetColumnNames = targetColumnNames;
+ this.inputColumnNames = inputColumnNames;
+ this.maxRetainedSize = maxRetainedSize;
+ this.maxReturnSize = maxReturnSize;
+ this.totalRow = 0;
+
+ if (modelInferenceDescriptor.getInferenceWindowParameter() != null) {
+ windowType = modelInferenceDescriptor.getInferenceWindowParameter().getWindowType();
+ }
+ }
+
+ @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 inferenceExecutionFuture;
+ } else if (executionDone) {
+ return childBlocked;
+ } else {
+ return successfulAsList(Arrays.asList(inferenceExecutionFuture, childBlocked));
+ }
+ }
+
+ private boolean forecastExecutionDone() {
+ if (inferenceExecutionFuture == null) {
+ return true;
+ }
+ return inferenceExecutionFuture.isDone();
+ }
+
+ @Override
+ public boolean hasNext() throws Exception {
+ return !finished || (results != null && results.size() != resultIndex);
+ }
+
+ @Override
+ public TsBlock next() throws Exception {
+ if (inferenceExecutionFuture == null) {
+ if (child.hasNextWithTimer()) {
+ TsBlock inputTsBlock = child.nextWithTimer();
+ if (inputTsBlock != null) {
+ appendTsBlockToBuilder(inputTsBlock);
+ }
+ } else {
+ submitInferenceTask();
+ }
+ return null;
+ } else {
+
+ if (results != null && resultIndex != results.size()) {
+ TsBlock tsBlock = serde.deserialize(results.get(resultIndex));
+ resultIndex++;
+ return tsBlock;
+ }
+
+ try {
+ if (!inferenceExecutionFuture.isDone()) {
+ throw new IllegalStateException(
+ "The operator cannot continue until the forecast execution is done.");
+ }
+
+ TInferenceResp inferenceResp = inferenceExecutionFuture.get();
+ if (inferenceResp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ String message =
+ String.format(
+ "Error occurred while executing inference:[%s]",
+ inferenceResp.getStatus().getMessage());
+ throw new ModelInferenceProcessException(message);
+ }
+
+ finished = true;
+ TsBlock resultTsBlock = serde.deserialize(inferenceResp.inferenceResult.get(0));
+ results = inferenceResp.inferenceResult;
+ resultIndex++;
+ return resultTsBlock;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new ModelInferenceProcessException(e.getMessage());
+ } catch (ExecutionException e) {
+ throw new ModelInferenceProcessException(e.getMessage());
+ }
+ }
+ }
+
+ private void appendTsBlockToBuilder(TsBlock inputTsBlock) {
+ TimeColumnBuilder timeColumnBuilder = inputTsBlockBuilder.getTimeColumnBuilder();
+ ColumnBuilder[] columnBuilders = inputTsBlockBuilder.getValueColumnBuilders();
+ totalRow += inputTsBlock.getPositionCount();
+ 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 TWindowParams getWindowParams() {
+ TWindowParams windowParams;
+ if (windowType == null) {
+ return null;
+ }
+ if (windowType == InferenceWindowType.COUNT) {
+ CountInferenceWindowParameter countInferenceWindowParameter =
+ (CountInferenceWindowParameter) modelInferenceDescriptor.getInferenceWindowParameter();
+ windowParams = new TWindowParams();
+ windowParams.setWindowInterval((int) countInferenceWindowParameter.getInterval());
+ windowParams.setWindowStep((int) countInferenceWindowParameter.getStep());
+ } else {
+ windowParams = null;
+ }
+ return windowParams;
+ }
+
+ private TsBlock preProcess(TsBlock inputTsBlock) {
+ boolean notBuiltIn = !modelInferenceDescriptor.getModelInformation().isBuiltIn();
+ if (windowType == null || windowType == InferenceWindowType.HEAD) {
+ if (notBuiltIn
+ && totalRow != modelInferenceDescriptor.getModelInformation().getInputShape()[0]) {
+ throw new ModelInferenceProcessException(
+ String.format(
+ "The number of rows %s in the input data does not match the model input %s. Try to use LIMIT in SQL or WINDOW in CALL INFERENCE",
+ totalRow, modelInferenceDescriptor.getModelInformation().getInputShape()[0]));
+ }
+ return inputTsBlock;
+ } else if (windowType == InferenceWindowType.COUNT) {
+ if (notBuiltIn
+ && totalRow < modelInferenceDescriptor.getModelInformation().getInputShape()[0]) {
+ throw new ModelInferenceProcessException(
+ String.format(
+ "The number of rows %s in the input data is less than the model input %s. ",
+ totalRow, modelInferenceDescriptor.getModelInformation().getInputShape()[0]));
+ }
+ } else if (windowType == InferenceWindowType.TAIL) {
+ if (notBuiltIn
+ && totalRow < modelInferenceDescriptor.getModelInformation().getInputShape()[0]) {
+ throw new ModelInferenceProcessException(
+ String.format(
+ "The number of rows %s in the input data is less than the model input %s. ",
+ totalRow, modelInferenceDescriptor.getModelInformation().getInputShape()[0]));
+ }
+ // Tail window logic: get the latest data for inference
+ long windowSize =
+ (int)
+ ((BottomInferenceWindowParameter)
+ modelInferenceDescriptor.getInferenceWindowParameter())
+ .getWindowSize();
+ return inputTsBlock.subTsBlock((int) (totalRow - windowSize));
+ }
+ return inputTsBlock;
+ }
+
+ private void submitInferenceTask() {
+
+ TsBlock inputTsBlock = inputTsBlockBuilder.build();
+
+ TsBlock finalInputTsBlock = preProcess(inputTsBlock);
+ TWindowParams windowParams = getWindowParams();
+
+ Map columnNameIndexMap = new HashMap<>();
+
+ for (int i = 0; i < inputColumnNames.size(); i++) {
+ columnNameIndexMap.put(inputColumnNames.get(i), i);
+ }
+
+ inferenceExecutionFuture =
+ Futures.submit(
+ () -> {
+ try (AINodeClient client =
+ AINodeClientManager.getInstance()
+ .borrowClient(modelInferenceDescriptor.getTargetAINode())) {
+ return client.inference(
+ modelInferenceDescriptor.getModelName(),
+ targetColumnNames,
+ Arrays.stream(modelInferenceDescriptor.getModelInformation().getInputDataType())
+ .map(TSDataType::toString)
+ .collect(Collectors.toList()),
+ columnNameIndexMap,
+ finalInputTsBlock,
+ modelInferenceDescriptor.getInferenceAttributes(),
+ windowParams);
+ } catch (Exception e) {
+ throw new ModelInferenceProcessException(e.getMessage());
+ }
+ },
+ modelInferenceExecutor);
+ }
+
+ @Override
+ public boolean isFinished() throws Exception {
+ return finished && !hasNext();
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (inferenceExecutionFuture != null) {
+ inferenceExecutionFuture.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();
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return INSTANCE_SIZE
+ + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(child)
+ + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext)
+ + inputTsBlockBuilder.getRetainedSizeInBytes()
+ + (inputColumnNames == null
+ ? 0
+ : inputColumnNames.stream().mapToLong(RamUsageEstimator::sizeOf).sum())
+ + (targetColumnNames == null
+ ? 0
+ : targetColumnNames.stream().mapToLong(RamUsageEstimator::sizeOf).sum());
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/BottomInferenceWindowParameter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/BottomInferenceWindowParameter.java
new file mode 100644
index 000000000000..77953b122f4b
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/BottomInferenceWindowParameter.java
@@ -0,0 +1,72 @@
+/*
+ * 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.window.ainode;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class BottomInferenceWindowParameter extends InferenceWindowParameter {
+
+ long windowSize;
+
+ public BottomInferenceWindowParameter(long windowSize) {
+ this.windowSize = windowSize;
+ this.windowType = InferenceWindowType.TAIL;
+ }
+
+ public long getWindowSize() {
+ return windowSize;
+ }
+
+ @Override
+ public void serializeAttributes(ByteBuffer buffer) {
+ ReadWriteIOUtils.write(windowSize, buffer);
+ }
+
+ @Override
+ public void serializeAttributes(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(windowSize, stream);
+ }
+
+ public static BottomInferenceWindowParameter deserialize(ByteBuffer byteBuffer) {
+ long windowSize = byteBuffer.getLong();
+ return new BottomInferenceWindowParameter(windowSize);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (!(obj instanceof BottomInferenceWindowParameter)) {
+ return false;
+ }
+ BottomInferenceWindowParameter parameter = (BottomInferenceWindowParameter) obj;
+ return windowSize == parameter.windowSize;
+ }
+
+ @Override
+ public int hashCode() {
+ return Long.hashCode(windowSize);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindow.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindow.java
new file mode 100644
index 000000000000..723e87593464
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindow.java
@@ -0,0 +1,40 @@
+/*
+ * 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.window.ainode;
+
+public class CountInferenceWindow extends InferenceWindow {
+
+ private final long interval;
+ private final long step;
+
+ public CountInferenceWindow(long interval, long step) {
+ super(InferenceWindowType.COUNT);
+ this.interval = interval;
+ this.step = step;
+ }
+
+ public long getInterval() {
+ return interval;
+ }
+
+ public long getStep() {
+ return step;
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindowParameter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindowParameter.java
new file mode 100644
index 000000000000..6a6371c4a70a
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/CountInferenceWindowParameter.java
@@ -0,0 +1,82 @@
+/*
+ * 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.window.ainode;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+public class CountInferenceWindowParameter extends InferenceWindowParameter {
+
+ private final long interval;
+ private final long step;
+
+ public CountInferenceWindowParameter(long interval, long step) {
+ this.windowType = InferenceWindowType.COUNT;
+ this.interval = interval;
+ this.step = step;
+ }
+
+ public long getInterval() {
+ return interval;
+ }
+
+ public long getStep() {
+ return step;
+ }
+
+ @Override
+ public void serializeAttributes(ByteBuffer buffer) {
+ ReadWriteIOUtils.write(interval, buffer);
+ ReadWriteIOUtils.write(step, buffer);
+ }
+
+ @Override
+ public void serializeAttributes(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(interval, stream);
+ ReadWriteIOUtils.write(step, stream);
+ }
+
+ public static CountInferenceWindowParameter deserialize(ByteBuffer byteBuffer) {
+ long interval = ReadWriteIOUtils.readLong(byteBuffer);
+ long step = ReadWriteIOUtils.readLong(byteBuffer);
+ return new CountInferenceWindowParameter(interval, step);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (!(obj instanceof CountInferenceWindowParameter)) {
+ return false;
+ }
+ CountInferenceWindowParameter parameter = (CountInferenceWindowParameter) obj;
+ return interval == parameter.interval && step == parameter.step;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(interval, step);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/HeadInferenceWindow.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/HeadInferenceWindow.java
new file mode 100644
index 000000000000..8e4f2cc65cb8
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/HeadInferenceWindow.java
@@ -0,0 +1,33 @@
+/*
+ * 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.window.ainode;
+
+public class HeadInferenceWindow extends InferenceWindow {
+ private final long windowSize;
+
+ public HeadInferenceWindow(long windowSize) {
+ super(InferenceWindowType.HEAD);
+ this.windowSize = windowSize;
+ }
+
+ public long getWindowSize() {
+ return windowSize;
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindow.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindow.java
new file mode 100644
index 000000000000..e5c00f910d0c
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindow.java
@@ -0,0 +1,32 @@
+/*
+ * 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.window.ainode;
+
+public class InferenceWindow {
+ private final InferenceWindowType type;
+
+ public InferenceWindow(InferenceWindowType type) {
+ this.type = type;
+ }
+
+ public InferenceWindowType getType() {
+ return type;
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowParameter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowParameter.java
new file mode 100644
index 000000000000..b9ab1343c3b8
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowParameter.java
@@ -0,0 +1,63 @@
+/*
+ * 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.window.ainode;
+
+import org.apache.iotdb.db.exception.sql.SemanticException;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public abstract class InferenceWindowParameter {
+
+ protected InferenceWindowType windowType;
+
+ public InferenceWindowType getWindowType() {
+ return windowType;
+ }
+
+ public abstract void serializeAttributes(ByteBuffer buffer);
+
+ public abstract void serializeAttributes(DataOutputStream stream) throws IOException;
+
+ public void serialize(ByteBuffer buffer) {
+ ReadWriteIOUtils.write(windowType.ordinal(), buffer);
+ serializeAttributes(buffer);
+ }
+
+ public void serialize(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(windowType.ordinal(), stream);
+ serializeAttributes(stream);
+ }
+
+ public static InferenceWindowParameter deserialize(ByteBuffer byteBuffer) {
+ InferenceWindowType windowType =
+ InferenceWindowType.values()[ReadWriteIOUtils.readInt(byteBuffer)];
+ if (windowType == InferenceWindowType.TAIL) {
+ return BottomInferenceWindowParameter.deserialize(byteBuffer);
+ } else if (windowType == InferenceWindowType.COUNT) {
+ return CountInferenceWindowParameter.deserialize(byteBuffer);
+ } else {
+ throw new SemanticException("Unsupported inference window type: " + windowType);
+ }
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowType.java
new file mode 100644
index 000000000000..f792327396f9
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/InferenceWindowType.java
@@ -0,0 +1,25 @@
+/*
+ * 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.window.ainode;
+
+public enum InferenceWindowType {
+ HEAD,
+ TAIL,
+ COUNT
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/TailInferenceWindow.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/TailInferenceWindow.java
new file mode 100644
index 000000000000..3bbc568cac2b
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/window/ainode/TailInferenceWindow.java
@@ -0,0 +1,34 @@
+/*
+ * 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.window.ainode;
+
+public class TailInferenceWindow extends InferenceWindow {
+
+ private final long windowSize;
+
+ public TailInferenceWindow(long windowSize) {
+ super(InferenceWindowType.TAIL);
+ this.windowSize = windowSize;
+ }
+
+ public long getWindowSize() {
+ return windowSize;
+ }
+}
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 94a1e518b74b..93e386b9bd52 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
@@ -25,6 +25,7 @@
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.common.rpc.thrift.TSchemaNode;
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.commons.partition.DataPartition;
import org.apache.iotdb.commons.partition.SchemaPartition;
import org.apache.iotdb.commons.path.PartialPath;
@@ -49,6 +50,7 @@
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.StatementType;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
@@ -237,6 +239,8 @@ aggregation results last_value(temperature) and last_value(status), whereas buck
// indicate whether the Nodes produce source data are VirtualSourceNodes
private boolean isVirtualSource = false;
+ private ModelInferenceDescriptor modelInferenceDescriptor;
+
/////////////////////////////////////////////////////////////////////////////////////////////////
// SELECT INTO Analysis
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -892,6 +896,21 @@ public Map getOutputDeviceToQueriedDevicesMap() {
return outputDeviceToQueriedDevicesMap;
}
+ public ModelInferenceDescriptor getModelInferenceDescriptor() {
+ return modelInferenceDescriptor;
+ }
+
+ public ModelInformation getModelInformation() {
+ if (modelInferenceDescriptor == null) {
+ return null;
+ }
+ return modelInferenceDescriptor.getModelInformation();
+ }
+
+ public void setModelInferenceDescriptor(ModelInferenceDescriptor modelInferenceDescriptor) {
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ }
+
public void setOutputDeviceToQueriedDevicesMap(
Map outputDeviceToQueriedDevicesMap) {
this.outputDeviceToQueriedDevicesMap = 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 605fb4c95dcb..b01ee8babed1 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
@@ -20,12 +20,14 @@
package org.apache.iotdb.db.queryengine.plan.analyze;
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
import org.apache.iotdb.commons.client.exception.ClientManagerException;
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.ModelInformation;
import org.apache.iotdb.commons.partition.DataPartition;
import org.apache.iotdb.commons.partition.DataPartitionQueryParam;
import org.apache.iotdb.commons.partition.SchemaNodeManagementPartition;
@@ -41,6 +43,7 @@
import org.apache.iotdb.confignode.rpc.thrift.TGetDataNodeLocationsResp;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.exception.ainode.GetModelInfoException;
import org.apache.iotdb.db.exception.metadata.template.TemplateIncompatibleException;
import org.apache.iotdb.db.exception.metadata.view.UnsupportedViewException;
import org.apache.iotdb.db.exception.sql.SemanticException;
@@ -59,6 +62,14 @@
import org.apache.iotdb.db.queryengine.common.schematree.IMeasurementSchemaInfo;
import org.apache.iotdb.db.queryengine.common.schematree.ISchemaTree;
import org.apache.iotdb.db.queryengine.execution.operator.window.WindowType;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.BottomInferenceWindowParameter;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.CountInferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.CountInferenceWindowParameter;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.HeadInferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindowParameter;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindowType;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.TailInferenceWindow;
import org.apache.iotdb.db.queryengine.metric.QueryPlanCostMetricSet;
import org.apache.iotdb.db.queryengine.plan.analyze.lock.DataNodeSchemaLockManager;
import org.apache.iotdb.db.queryengine.plan.analyze.lock.SchemaLockType;
@@ -217,11 +228,14 @@ public class AnalyzeVisitor extends StatementVisitor
public static final Expression END_TIME_EXPRESSION =
TimeSeriesOperand.constructColumnHeaderExpression(ENDTIME, TSDataType.INT64);
+ private static final String INFERENCE_COLUMN_NAME = "output";
+
private final List lastQueryColumnNames =
new ArrayList<>(Arrays.asList("TIME", "TIMESERIES", "VALUE", "DATATYPE"));
private final IPartitionFetcher partitionFetcher;
private final ISchemaFetcher schemaFetcher;
+ private final IModelFetcher modelFetcher;
private static final PerformanceOverviewMetrics PERFORMANCE_OVERVIEW_METRICS =
PerformanceOverviewMetrics.getInstance();
@@ -229,6 +243,7 @@ public class AnalyzeVisitor extends StatementVisitor
public AnalyzeVisitor(IPartitionFetcher partitionFetcher, ISchemaFetcher schemaFetcher) {
this.partitionFetcher = partitionFetcher;
this.schemaFetcher = schemaFetcher;
+ this.modelFetcher = ModelFetcher.getInstance();
}
@Override
@@ -268,6 +283,9 @@ public Analysis visitQuery(QueryStatement queryStatement, MPPQueryContext contex
// check for semantic errors
queryStatement.semanticCheck();
+ // fetch model inference information and check
+ 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
@@ -397,6 +415,77 @@ public Analysis visitQuery(QueryStatement queryStatement, MPPQueryContext contex
return analysis;
}
+ // check if there is proper model to inference for MODEL_NAME, there is no need to do the
+ // following analyze if there isn't.
+ private void analyzeModelInference(Analysis analysis, QueryStatement queryStatement) {
+ if (!queryStatement.hasModelInference()) {
+ return;
+ }
+
+ // Get model metadata from configNode and do some check
+ String modelId = queryStatement.getModelName();
+ TSStatus status = modelFetcher.fetchModel(modelId, analysis);
+ if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ throw new GetModelInfoException(status.getMessage());
+ }
+ ModelInformation modelInformation = analysis.getModelInformation();
+ if (modelInformation == null || !modelInformation.available()) {
+ throw new SemanticException("Model " + modelId + " is not active");
+ }
+
+ // set inference window if there is
+ if (queryStatement.isSetInferenceWindow()) {
+ InferenceWindow window = queryStatement.getInferenceWindow();
+ if (InferenceWindowType.HEAD == window.getType()) {
+ long windowSize = ((HeadInferenceWindow) window).getWindowSize();
+ checkWindowSize(windowSize, modelInformation);
+ if (queryStatement.hasLimit() && queryStatement.getRowLimit() < windowSize) {
+ throw new SemanticException(
+ "Limit in Sql should be larger than window size in inference");
+ }
+ // optimize head window by limitNode
+ queryStatement.setRowLimit(windowSize);
+ } else if (InferenceWindowType.TAIL == window.getType()) {
+ long windowSize = ((TailInferenceWindow) window).getWindowSize();
+ checkWindowSize(windowSize, modelInformation);
+ InferenceWindowParameter inferenceWindowParameter =
+ new BottomInferenceWindowParameter(windowSize);
+ analysis
+ .getModelInferenceDescriptor()
+ .setInferenceWindowParameter(inferenceWindowParameter);
+ } else if (InferenceWindowType.COUNT == window.getType()) {
+ CountInferenceWindow countInferenceWindow = (CountInferenceWindow) window;
+ checkWindowSize(countInferenceWindow.getInterval(), modelInformation);
+ InferenceWindowParameter inferenceWindowParameter =
+ new CountInferenceWindowParameter(
+ countInferenceWindow.getInterval(), countInferenceWindow.getStep());
+ analysis
+ .getModelInferenceDescriptor()
+ .setInferenceWindowParameter(inferenceWindowParameter);
+ }
+ }
+
+ // set inference attributes if there is
+ if (queryStatement.hasInferenceAttributes()) {
+ analysis
+ .getModelInferenceDescriptor()
+ .setInferenceAttributes(queryStatement.getInferenceAttributes());
+ }
+ }
+
+ private void checkWindowSize(long windowSize, ModelInformation modelInformation) {
+ if (modelInformation.isBuiltIn()) {
+ return;
+ }
+
+ if (modelInformation.getInputShape()[0] != windowSize) {
+ throw new SemanticException(
+ String.format(
+ "Window output %d is not equal to input size of model %d",
+ windowSize, modelInformation.getInputShape()[0]));
+ }
+ }
+
private ISchemaTree analyzeSchema(
QueryStatement queryStatement, Analysis analysis, MPPQueryContext context) {
// concat path and construct path pattern tree
@@ -1574,6 +1663,27 @@ static void analyzeOutput(
return;
}
+ if (queryStatement.hasModelInference()) {
+ ModelInformation modelInformation = analysis.getModelInformation();
+ // check input
+ checkInputShape(modelInformation, outputExpressions);
+ checkInputType(analysis, modelInformation, outputExpressions);
+
+ // set output
+ List columnHeaders = new ArrayList<>();
+ int[] outputShape = modelInformation.getOutputShape();
+ TSDataType[] outputDataType = modelInformation.getOutputDataType();
+ for (int i = 0; i < outputShape[1]; i++) {
+ columnHeaders.add(new ColumnHeader(INFERENCE_COLUMN_NAME + i, outputDataType[i]));
+ }
+ analysis
+ .getModelInferenceDescriptor()
+ .setOutputColumnNames(
+ columnHeaders.stream().map(ColumnHeader::getColumnName).collect(Collectors.toList()));
+ analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, true));
+ return;
+ }
+
boolean isIgnoreTimestamp = queryStatement.isAggregationQuery() && !queryStatement.isGroupBy();
List columnHeaders = new ArrayList<>();
if (queryStatement.isAlignByDevice()) {
@@ -1592,6 +1702,72 @@ static void analyzeOutput(
analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, isIgnoreTimestamp));
}
+ // check if the result of SQL matches the input of model
+ private static void checkInputShape(
+ ModelInformation modelInformation, List> outputExpressions) {
+ if (modelInformation.isBuiltIn()) {
+ modelInformation.setInputColumnSize(outputExpressions.size());
+ return;
+ }
+
+ // check inputShape
+ int[] inputShape = modelInformation.getInputShape();
+ if (inputShape.length != 2) {
+ throw new SemanticException(
+ String.format(
+ "The input shape of model is not correct, the dimension of input shape should be 2, actual dimension is %d",
+ inputShape.length));
+ }
+ int columnNumber = inputShape[1];
+ if (columnNumber != outputExpressions.size()) {
+ throw new SemanticException(
+ String.format(
+ "The column number of SQL result does not match the number of model input [%d] for inference",
+ columnNumber));
+ }
+ }
+
+ private static void checkInputType(
+ Analysis analysis,
+ ModelInformation modelInformation,
+ List> outputExpressions) {
+
+ if (modelInformation.isBuiltIn()) {
+ TSDataType[] inputType = new TSDataType[outputExpressions.size()];
+ for (int i = 0; i < outputExpressions.size(); i++) {
+ Expression inputExpression = outputExpressions.get(i).left;
+ TSDataType inputDataType = analysis.getType(inputExpression);
+ if (!inputDataType.isNumeric()) {
+ throw new SemanticException(
+ String.format(
+ "The type of SQL result column [%s in %d] should be numeric when inference",
+ inputDataType, i));
+ }
+ inputType[i] = inputDataType;
+ }
+ modelInformation.setInputDataType(inputType);
+ return;
+ }
+
+ TSDataType[] inputType = modelInformation.getInputDataType();
+ if (inputType.length != modelInformation.getInputShape()[1]) {
+ throw new SemanticException(
+ String.format(
+ "The inputType does not match the input shape [%d] for inference",
+ modelInformation.getInputShape()[1]));
+ }
+ for (int i = 0; i < inputType.length; i++) {
+ Expression inputExpression = outputExpressions.get(i).left;
+ TSDataType inputDataType = analysis.getType(inputExpression);
+ if (inputDataType != inputType[i]) {
+ throw new SemanticException(
+ String.format(
+ "The type of SQL result column [%s in %d] does not match the type of model input [%s] when inference",
+ inputDataType, i, inputType[i]));
+ }
+ }
+ }
+
// For last query
private void analyzeLastOrderBy(Analysis analysis, QueryStatement queryStatement) {
if (!queryStatement.hasOrderBy()) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java
new file mode 100644
index 000000000000..1feecaefde9c
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IModelFetcher.java
@@ -0,0 +1,27 @@
+/*
+ * 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.common.rpc.thrift.TSStatus;
+
+public interface IModelFetcher {
+ /** Get model information by model id from configNode. */
+ TSStatus fetchModel(String modelId, Analysis analysis);
+}
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 000000000000..8cefb5e0cf3c
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ModelFetcher.java
@@ -0,0 +1,81 @@
+/*
+ * 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.common.rpc.thrift.TSStatus;
+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.ainode.ModelNotFoundException;
+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.db.queryengine.plan.planner.plan.parameter.model.ModelInferenceDescriptor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.apache.thrift.TException;
+
+public class ModelFetcher implements IModelFetcher {
+
+ private final IClientManager configNodeClientManager =
+ ConfigNodeClientManager.getInstance();
+
+ private static final class ModelFetcherHolder {
+
+ private static final ModelFetcher INSTANCE = new ModelFetcher();
+
+ private ModelFetcherHolder() {}
+ }
+
+ public static ModelFetcher getInstance() {
+ return ModelFetcherHolder.INSTANCE;
+ }
+
+ private ModelFetcher() {}
+
+ @Override
+ public TSStatus fetchModel(String modelName, Analysis analysis) {
+ try (ConfigNodeClient client =
+ configNodeClientManager.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TGetModelInfoResp getModelInfoResp = client.getModelInfo(new TGetModelInfoReq(modelName));
+ if (getModelInfoResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ if (getModelInfoResp.modelInfo != null && getModelInfoResp.isSetAiNodeAddress()) {
+ analysis.setModelInferenceDescriptor(
+ new ModelInferenceDescriptor(
+ getModelInfoResp.aiNodeAddress,
+ ModelInformation.deserialize(getModelInfoResp.modelInfo)));
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ } else {
+ TSStatus status = new TSStatus(TSStatusCode.GET_MODEL_INFO_ERROR.getStatusCode());
+ status.setMessage(String.format("model [%s] is not available", modelName));
+ return status;
+ }
+ } else {
+ throw new ModelNotFoundException(getModelInfoResp.getStatus().getMessage());
+ }
+ } catch (ClientManagerException | TException e) {
+ throw new StatementAnalyzeException(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 6a04b32c0ae6..b928d4722344 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
@@ -52,6 +52,10 @@
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.UnSetTTLTask;
+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.ShowAINodesTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowModelsTask;
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;
@@ -120,6 +124,10 @@
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowTriggersStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowVariablesStatement;
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.ShowAINodesStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -585,4 +593,29 @@ public IConfigTask visitShowThrottleQuota(
ShowThrottleQuotaStatement showThrottleQuotaStatement, MPPQueryContext context) {
return new ShowThrottleQuotaTask(showThrottleQuotaStatement);
}
+
+ /** AI Model Management */
+ @Override
+ public IConfigTask visitCreateModel(
+ CreateModelStatement createModelStatement, MPPQueryContext context) {
+ return new CreateModelTask(createModelStatement, context);
+ }
+
+ @Override
+ public IConfigTask visitDropModel(
+ DropModelStatement dropModelStatement, MPPQueryContext context) {
+ return new DropModelTask(dropModelStatement.getModelName());
+ }
+
+ @Override
+ public IConfigTask visitShowModels(
+ ShowModelsStatement showModelsStatement, MPPQueryContext context) {
+ return new ShowModelsTask(showModelsStatement.getModelName());
+ }
+
+ @Override
+ public IConfigTask visitShowAINodes(
+ ShowAINodesStatement showAINodesStatement, MPPQueryContext context) {
+ return new ShowAINodesTask(showAINodesStatement);
+ }
}
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 8b6d53582ebe..7ac6b553e5ec 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
@@ -63,6 +63,7 @@
import org.apache.iotdb.confignode.rpc.thrift.TCountTimeSlotListResp;
import org.apache.iotdb.confignode.rpc.thrift.TCreateCQReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreateFunctionReq;
+import org.apache.iotdb.confignode.rpc.thrift.TCreateModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreatePipePluginReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TCreateTopicReq;
@@ -74,6 +75,7 @@
import org.apache.iotdb.confignode.rpc.thrift.TDeleteTimeSeriesReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropCQReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropFunctionReq;
+import org.apache.iotdb.confignode.rpc.thrift.TDropModelReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropPipePluginReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TDropTopicReq;
@@ -92,11 +94,14 @@
import org.apache.iotdb.confignode.rpc.thrift.TPipeConfigTransferReq;
import org.apache.iotdb.confignode.rpc.thrift.TPipeConfigTransferResp;
import org.apache.iotdb.confignode.rpc.thrift.TRegionInfo;
+import org.apache.iotdb.confignode.rpc.thrift.TShowAINodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowCQResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowConfigNodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowDatabaseResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeInfo;
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
import org.apache.iotdb.confignode.rpc.thrift.TShowRegionReq;
@@ -150,6 +155,8 @@
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowTTLTask;
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.ShowAINodesTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model.ShowModelsTask;
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;
@@ -180,6 +187,8 @@
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowDatabaseStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowRegionStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowTTLStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAINodesStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -2645,6 +2654,90 @@ public SettableFuture showContinuousQueries() {
return future;
}
+ @Override
+ public SettableFuture createModel(
+ CreateModelStatement createModelStatement, MPPQueryContext context) {
+ SettableFuture future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TCreateModelReq req =
+ new TCreateModelReq(createModelStatement.getModelName(), createModelStatement.getUri());
+ final TSStatus status = client.createModel(req);
+ if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != status.getCode()) {
+ LOGGER.warn(
+ "[{}] Failed to create model {}. TSStatus is {}",
+ status,
+ createModelStatement.getModelName(),
+ status.message);
+ future.setException(new IoTDBException(status.message, status.code));
+ } else {
+ future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
+ }
+ } catch (ClientManagerException | TException e) {
+ future.setException(e);
+ }
+ return future;
+ }
+
+ @Override
+ public SettableFuture dropModel(String modelName) {
+ SettableFuture future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ final TSStatus executionStatus = client.dropModel(new TDropModelReq(modelName));
+ if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != executionStatus.getCode()) {
+ LOGGER.warn("[{}] Failed to drop model {}.", executionStatus, modelName);
+ 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 showModels(String modelName) {
+ SettableFuture future = SettableFuture.create();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ TShowModelReq req = new TShowModelReq();
+ if (modelName != null) {
+ req.setModelId(modelName);
+ }
+ TShowModelResp showModelResp = client.showModel(req);
+ 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 showAINodes(ShowAINodesStatement showAINodesStatement) {
+ SettableFuture future = SettableFuture.create();
+ TShowAINodesResp resp = new TShowAINodesResp();
+ try (ConfigNodeClient client =
+ CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ resp = client.showAINodes();
+ if (resp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ future.setException(new IoTDBException(resp.getStatus().message, resp.getStatus().code));
+ return future;
+ }
+ } catch (ClientManagerException | TException e) {
+ future.setException(e);
+ }
+ ShowAINodesTask.buildTsBlock(resp, future);
+ return future;
+ }
+
@Override
public SettableFuture setSpaceQuota(
SetSpaceQuotaStatement setSpaceQuotaStatement) {
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 fad6ec99cba8..b3911f559795 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
@@ -46,6 +46,8 @@
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowDatabaseStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowRegionStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowTTLStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAINodesStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -242,6 +244,15 @@ SettableFuture showThrottleQuota(
TThrottleQuotaResp getThrottleQuota();
+ SettableFuture createModel(
+ CreateModelStatement createModelStatement, MPPQueryContext context);
+
+ SettableFuture dropModel(String modelName);
+
+ SettableFuture showModels(String modelName);
+
+ SettableFuture showAINodes(ShowAINodesStatement showAINodesStatement);
+
TPipeTransferResp handleTransferConfigPlan(String clientId, TPipeTransferReq req);
void handlePipeConfigClientExit(String clientId);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterDetailsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterDetailsTask.java
index 54f6c5b57be0..f57c84fe3b93 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterDetailsTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterDetailsTask.java
@@ -41,6 +41,7 @@
import java.util.List;
import java.util.stream.Collectors;
+import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_AI_NODE;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_CONFIG_NODE;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_DATA_NODE;
@@ -104,6 +105,56 @@ private static void buildConfigNodesTsBlock(
builder.declarePosition();
}
+ private static void buildAINodeTsBlock(
+ TsBlockBuilder builder,
+ int nodeId,
+ String nodeStatus,
+ String internalAddress,
+ int internalPort,
+ TNodeVersionInfo versionInfo) {
+
+ builder.getTimeColumnBuilder().writeLong(0L);
+ builder.getColumnBuilder(0).writeInt(nodeId);
+ builder
+ .getColumnBuilder(1)
+ .writeBinary(new Binary(NODE_TYPE_AI_NODE, TSFileConfig.STRING_CHARSET));
+ if (nodeStatus == null) {
+ builder.getColumnBuilder(2).appendNull();
+ } else {
+ builder.getColumnBuilder(2).writeBinary(new Binary(nodeStatus, TSFileConfig.STRING_CHARSET));
+ }
+
+ if (internalAddress == null) {
+ builder.getColumnBuilder(3).appendNull();
+ } else {
+ builder
+ .getColumnBuilder(3)
+ .writeBinary(new Binary(internalAddress, TSFileConfig.STRING_CHARSET));
+ }
+ builder.getColumnBuilder(4).writeInt(internalPort);
+ builder.getColumnBuilder(5).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ builder.getColumnBuilder(6).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ builder.getColumnBuilder(7).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ builder.getColumnBuilder(8).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ builder.getColumnBuilder(9).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ builder.getColumnBuilder(10).writeBinary(new Binary("", TSFileConfig.STRING_CHARSET));
+ if (versionInfo == null || versionInfo.getVersion() == null) {
+ builder.getColumnBuilder(11).appendNull();
+ } else {
+ builder
+ .getColumnBuilder(11)
+ .writeBinary(new Binary(versionInfo.getVersion(), TSFileConfig.STRING_CHARSET));
+ }
+ if (versionInfo == null || versionInfo.getBuildInfo() == null) {
+ builder.getColumnBuilder(12).appendNull();
+ } else {
+ builder
+ .getColumnBuilder(12)
+ .writeBinary(new Binary(versionInfo.getBuildInfo(), TSFileConfig.STRING_CHARSET));
+ }
+ builder.declarePosition();
+ }
+
@SuppressWarnings("squid:S107")
private static void buildDataNodesTsBlock(
TsBlockBuilder builder,
@@ -208,6 +259,17 @@ public static void buildTSBlock(
e.getSchemaRegionConsensusEndPoint().getPort(),
e.getDataRegionConsensusEndPoint().getPort(),
clusterNodeInfos.getNodeVersionInfo().get(e.getDataNodeId())));
+ clusterNodeInfos
+ .getAiNodeList()
+ .forEach(
+ e ->
+ buildAINodeTsBlock(
+ builder,
+ e.getAiNodeId(),
+ clusterNodeInfos.getNodeStatus().get(e.getAiNodeId()),
+ e.getInternalEndPoint().getIp(),
+ e.getInternalEndPoint().getPort(),
+ clusterNodeInfos.getNodeVersionInfo().get(e.getAiNodeId())));
DatasetHeader datasetHeader = DatasetHeaderFactory.getShowClusterDetailsHeader();
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/execution/config/metadata/ShowClusterTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterTask.java
index 7d0c41f6baad..960b088293c5 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowClusterTask.java
@@ -41,6 +41,7 @@
import java.util.List;
import java.util.stream.Collectors;
+import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_AI_NODE;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_CONFIG_NODE;
import static org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant.NODE_TYPE_DATA_NODE;
@@ -136,6 +137,20 @@ public static void buildTsBlock(
e.getInternalEndPoint().getPort(),
clusterNodeInfos.getNodeVersionInfo().get(e.getDataNodeId())));
+ if (clusterNodeInfos.getAiNodeList() != null) {
+ clusterNodeInfos
+ .getAiNodeList()
+ .forEach(
+ e ->
+ buildTsBlock(
+ builder,
+ e.getAiNodeId(),
+ NODE_TYPE_AI_NODE,
+ clusterNodeInfos.getNodeStatus().get(e.getAiNodeId()),
+ e.getInternalEndPoint().getIp(),
+ e.getInternalEndPoint().getPort(),
+ clusterNodeInfos.getNodeVersionInfo().get(e.getAiNodeId())));
+ }
DatasetHeader datasetHeader = DatasetHeaderFactory.getShowClusterHeader();
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/execution/config/metadata/model/CreateModelTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/CreateModelTask.java
new file mode 100644
index 000000000000..842b558fd8a0
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/CreateModelTask.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.execution.config.metadata.model;
+
+import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
+import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
+import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+public class CreateModelTask implements IConfigTask {
+
+ private final CreateModelStatement createModelStatement;
+ private final MPPQueryContext context;
+
+ public CreateModelTask(CreateModelStatement createModelStatement, MPPQueryContext context) {
+ this.createModelStatement = createModelStatement;
+ this.context = context;
+ }
+
+ @Override
+ public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor)
+ throws InterruptedException {
+ return configTaskExecutor.createModel(createModelStatement, context);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/DropModelTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/DropModelTask.java
new file mode 100644
index 000000000000..f8db88790d4d
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/DropModelTask.java
@@ -0,0 +1,41 @@
+/*
+ * 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.execution.config.metadata.model;
+
+import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
+import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+public class DropModelTask implements IConfigTask {
+
+ private final String modelName;
+
+ public DropModelTask(String modelName) {
+ this.modelName = modelName;
+ }
+
+ @Override
+ public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor)
+ throws InterruptedException {
+ return configTaskExecutor.dropModel(modelName);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowAINodesTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowAINodesTask.java
new file mode 100644
index 000000000000..f63a43aa1a36
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowAINodesTask.java
@@ -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.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.execution.config.metadata.model;
+
+import org.apache.iotdb.confignode.rpc.thrift.TAINodeInfo;
+import org.apache.iotdb.confignode.rpc.thrift.TShowAINodesResp;
+import org.apache.iotdb.db.queryengine.common.header.ColumnHeader;
+import org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory;
+import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
+import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAINodesStatement;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.utils.BytesUtils;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class ShowAINodesTask implements IConfigTask {
+
+ private final ShowAINodesStatement showAINodesStatement;
+
+ public ShowAINodesTask(ShowAINodesStatement showAINodesStatement) {
+ this.showAINodesStatement = showAINodesStatement;
+ }
+
+ @Override
+ public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor)
+ throws InterruptedException {
+ return configTaskExecutor.showAINodes(showAINodesStatement);
+ }
+
+ public static void buildTsBlock(
+ TShowAINodesResp showAINodesResp, SettableFuture future) {
+ List outputDataTypes =
+ ColumnHeaderConstant.showAINodesColumnHeaders.stream()
+ .map(ColumnHeader::getColumnType)
+ .collect(Collectors.toList());
+ TsBlockBuilder builder = new TsBlockBuilder(outputDataTypes);
+ if (showAINodesResp.getAiNodesInfoList() != null) {
+ for (TAINodeInfo aiNodeInfo : showAINodesResp.getAiNodesInfoList()) {
+ builder.getTimeColumnBuilder().writeLong(0);
+ builder.getColumnBuilder(0).writeInt(aiNodeInfo.getAiNodeId());
+ builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(aiNodeInfo.getStatus()));
+ builder
+ .getColumnBuilder(2)
+ .writeBinary(BytesUtils.valueOf(aiNodeInfo.getInternalAddress()));
+ builder.getColumnBuilder(3).writeInt(aiNodeInfo.getInternalPort());
+
+ builder.declarePosition();
+ }
+ }
+ DatasetHeader datasetHeader = DatasetHeaderFactory.getShowAINodesHeader();
+ 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/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
new file mode 100644
index 000000000000..0dc70f21e3ee
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/model/ShowModelsTask.java
@@ -0,0 +1,112 @@
+/*
+ * 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.execution.config.metadata.model;
+
+import org.apache.iotdb.commons.model.ModelType;
+import org.apache.iotdb.db.queryengine.common.header.ColumnHeader;
+import org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory;
+import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
+import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.utils.BytesUtils;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class ShowModelsTask implements IConfigTask {
+
+ private final String modelName;
+
+ public ShowModelsTask(String modelName) {
+ this.modelName = modelName;
+ }
+
+ private static final String INPUT_SHAPE = "inputShape:";
+ private static final String OUTPUT_SHAPE = "outputShape:";
+ private static final String INPUT_DATA_TYPE = "inputDataType:";
+ private static final String OUTPUT_DATA_TYPE = "outputDataType:";
+ private static final String EMPTY_STRING = "";
+
+ @Override
+ public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor)
+ throws InterruptedException {
+ return configTaskExecutor.showModels(modelName);
+ }
+
+ public static void buildTsBlock(
+ List modelInfoList, SettableFuture future) {
+ List outputDataTypes =
+ ColumnHeaderConstant.showModelsColumnHeaders.stream()
+ .map(ColumnHeader::getColumnType)
+ .collect(Collectors.toList());
+ TsBlockBuilder builder = new TsBlockBuilder(outputDataTypes);
+ for (ByteBuffer modelInfo : modelInfoList) {
+ String modelId = ReadWriteIOUtils.readString(modelInfo);
+ String modelType = ReadWriteIOUtils.readString(modelInfo);
+ String state = ReadWriteIOUtils.readString(modelInfo);
+ String note;
+ String config;
+ if (Objects.equals(modelType, ModelType.USER_DEFINED.toString())) {
+ String inputShape = ReadWriteIOUtils.readString(modelInfo);
+ String outputShape = ReadWriteIOUtils.readString(modelInfo);
+ String inputTypes = ReadWriteIOUtils.readString(modelInfo);
+ String outputTypes = ReadWriteIOUtils.readString(modelInfo);
+ note = ReadWriteIOUtils.readString(modelInfo);
+ config =
+ INPUT_SHAPE
+ + inputShape
+ + OUTPUT_SHAPE
+ + outputShape
+ + INPUT_DATA_TYPE
+ + inputTypes
+ + OUTPUT_DATA_TYPE
+ + outputTypes;
+ } else {
+ config = EMPTY_STRING;
+ note = "Built-in model in IoTDB";
+ }
+
+ builder.getTimeColumnBuilder().writeLong(0L);
+ builder.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(modelId));
+ builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(modelType));
+ builder.getColumnBuilder(2).writeBinary(BytesUtils.valueOf(state));
+ builder.getColumnBuilder(3).writeBinary(BytesUtils.valueOf(config));
+ if (note != null) {
+ builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf(note));
+ } else {
+ builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf(""));
+ }
+ builder.declarePosition();
+ }
+ DatasetHeader datasetHeader = DatasetHeaderFactory.getShowModelsHeader();
+ 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/parser/ASTVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
index 7235ffbc5185..099678ffcee2 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
@@ -53,6 +53,10 @@
import org.apache.iotdb.db.qp.sql.IoTDBSqlParserBaseVisitor;
import org.apache.iotdb.db.queryengine.common.header.ColumnHeaderConstant;
import org.apache.iotdb.db.queryengine.execution.operator.window.WindowType;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.CountInferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.HeadInferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindow;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.TailInferenceWindow;
import org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer;
import org.apache.iotdb.db.queryengine.plan.expression.Expression;
import org.apache.iotdb.db.queryengine.plan.expression.ExpressionType;
@@ -161,6 +165,10 @@
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowTriggersStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowVariablesStatement;
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.ShowAINodesStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -1300,6 +1308,43 @@ private void parseViewSourcePaths(
}
}
+ // Create Model =====================================================================
+ public static void validateModelName(String modelName) {
+ if (modelName.length() < 2 || modelName.length() > 64) {
+ throw new SemanticException("Model name should be 2-64 characters");
+ } else if (modelName.startsWith("_")) {
+ throw new SemanticException("Model name should not start with '_'");
+ } else if (!modelName.matches("^[-\\w]*$")) {
+ throw new SemanticException("ModelName can only contain letters, numbers, and underscores");
+ }
+ }
+
+ @Override
+ public Statement visitCreateModel(IoTDBSqlParser.CreateModelContext ctx) {
+ CreateModelStatement createModelStatement = new CreateModelStatement();
+ String modelName = ctx.modelName.getText();
+ validateModelName(modelName);
+ createModelStatement.setModelName(parseIdentifier(modelName));
+ createModelStatement.setUri(ctx.modelUri.getText());
+ return createModelStatement;
+ }
+
+ // Drop Model =====================================================================
+ @Override
+ public Statement visitDropModel(IoTDBSqlParser.DropModelContext ctx) {
+ return new DropModelStatement(parseIdentifier(ctx.modelId.getText()));
+ }
+
+ // Show Models =====================================================================
+ @Override
+ public Statement visitShowModels(IoTDBSqlParser.ShowModelsContext ctx) {
+ ShowModelsStatement statement = new ShowModelsStatement();
+ if (ctx.modelId != null) {
+ statement.setModelName(parseIdentifier(ctx.modelId.getText()));
+ }
+ return statement;
+ }
+
/** Data Manipulation Language (DML). */
// Select Statement ========================================================================
@@ -3426,6 +3471,11 @@ public Statement visitShowConfigNodes(IoTDBSqlParser.ShowConfigNodesContext ctx)
return new ShowConfigNodesStatement();
}
+ @Override
+ public Statement visitShowAINodes(IoTDBSqlParser.ShowAINodesContext ctx) {
+ return new ShowAINodesStatement();
+ }
+
// device template
@Override
@@ -4320,6 +4370,57 @@ public Statement visitShowSpaceQuota(IoTDBSqlParser.ShowSpaceQuotaContext ctx) {
return showSpaceQuotaStatement;
}
+ @Override
+ public Statement visitCallInference(IoTDBSqlParser.CallInferenceContext ctx) {
+ String sql = ctx.inputSql.getText();
+ QueryStatement statement =
+ (QueryStatement)
+ StatementGenerator.createStatement(sql.substring(1, sql.length() - 1), zoneId);
+
+ statement.setModelName(parseIdentifier(ctx.modelId.getText()));
+ statement.setHasModelInference(true);
+
+ if (ctx.hparamPair() != null) {
+ for (IoTDBSqlParser.HparamPairContext context : ctx.hparamPair()) {
+ IoTDBSqlParser.HparamValueContext valueContext = context.hparamValue();
+ String paramKey = context.hparamKey.getText();
+ if (paramKey.equalsIgnoreCase("WINDOW")) {
+ if (statement.isSetInferenceWindow()) {
+ throw new SemanticException("There should be only one window in CALL INFERENCE.");
+ }
+ if (valueContext.windowFunction().isEmpty()) {
+ throw new SemanticException(
+ "Window Function(e.g. HEAD, TAIL, COUNT) should be set in value when key is 'WINDOW' in CALL INFERENCE");
+ }
+ parseWindowFunctionInInference(valueContext.windowFunction(), statement);
+ } else {
+ statement.addInferenceAttribute(
+ paramKey, parseAttributeValue(valueContext.attributeValue()));
+ }
+ }
+ }
+
+ return statement;
+ }
+
+ private void parseWindowFunctionInInference(
+ IoTDBSqlParser.WindowFunctionContext windowContext, QueryStatement statement) {
+ InferenceWindow inferenceWindow = null;
+ if (windowContext.TAIL() != null) {
+ inferenceWindow =
+ new TailInferenceWindow(Integer.parseInt(windowContext.windowSize.getText()));
+ } else if (windowContext.HEAD() != null) {
+ inferenceWindow =
+ new HeadInferenceWindow(Integer.parseInt(windowContext.windowSize.getText()));
+ } else if (windowContext.COUNT() != null) {
+ inferenceWindow =
+ new CountInferenceWindow(
+ Integer.parseInt(windowContext.interval.getText()),
+ Integer.parseInt(windowContext.step.getText()));
+ }
+ statement.setInferenceWindow(inferenceWindow);
+ }
+
@Override
public Statement visitShowCurrentTimestamp(IoTDBSqlParser.ShowCurrentTimestampContext ctx) {
return new ShowCurrentTimestampStatement();
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 08cddf94ac6e..b220309df583 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
@@ -68,6 +68,7 @@
import org.apache.iotdb.db.utils.QueryDataSetUtils;
import org.apache.iotdb.db.utils.TimestampPrecisionUtils;
import org.apache.iotdb.db.utils.constant.SqlConstant;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchTimeseriesReq;
import org.apache.iotdb.service.rpc.thrift.TSAggregationQueryReq;
import org.apache.iotdb.service.rpc.thrift.TSCreateAlignedTimeseriesReq;
import org.apache.iotdb.service.rpc.thrift.TSCreateMultiTimeseriesReq;
@@ -861,4 +862,8 @@ private static PartialPath parseDatabaseRawString(String database) throws Illega
MetaFormatUtils.checkDatabase(database);
return databasePath;
}
+
+ public static Statement createStatement(TFetchTimeseriesReq fetchTimeseriesReq, ZoneId zoneId) {
+ return invokeParser(fetchTimeseriesReq.getQueryBody(), zoneId);
+ }
}
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 1af2387b0d51..94332ea7871f 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
@@ -59,6 +59,7 @@
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.SeriesSchemaFetchScanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.TimeSeriesCountNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.TimeSeriesSchemaScanNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AI.InferenceNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ColumnInjectNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.DeviceViewIntoNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.DeviceViewNode;
@@ -1369,6 +1370,19 @@ public LogicalPlanBuilder planOrderBy(QueryStatement queryStatement, Analysis an
return this;
}
+ public LogicalPlanBuilder planInference(Analysis analysis) {
+ this.root =
+ new InferenceNode(
+ context.getQueryId().genPlanNodeId(),
+ root,
+ analysis.getModelInferenceDescriptor(),
+ analysis.getOutputExpressions().stream()
+ .map(expressionStringPair -> expressionStringPair.left.getExpressionString())
+ .collect(Collectors.toList()));
+
+ return this;
+ }
+
public LogicalPlanBuilder planEndTimeColumnInject(
GroupByTimeParameter groupByTimeParameter, boolean ascending) {
this.root =
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 c179db107031..5913bc9f6658 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
@@ -229,6 +229,10 @@ public PlanNode visitQuery(QueryStatement queryStatement, MPPQueryContext contex
planBuilder = planBuilder.planLimit(queryStatement.getRowLimit());
}
+ if (queryStatement.hasModelInference()) {
+ planBuilder.planInference(analysis);
+ }
+
// 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 05425c23b2b3..662db079241c 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
@@ -21,6 +21,7 @@
import org.apache.iotdb.common.rpc.thrift.TAggregationType;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.model.ModelInformation;
import org.apache.iotdb.commons.path.AlignedPath;
import org.apache.iotdb.commons.path.MeasurementPath;
import org.apache.iotdb.commons.path.PartialPath;
@@ -185,6 +186,7 @@
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.SeriesSchemaFetchScanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.TimeSeriesCountNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.read.TimeSeriesSchemaScanNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AI.InferenceNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ActiveRegionScanMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationMergeSortNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationNode;
@@ -242,6 +244,7 @@
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.ModelInferenceDescriptor;
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;
@@ -268,6 +271,7 @@
import org.apache.tsfile.file.metadata.PlainDeviceID;
import org.apache.tsfile.read.TimeValuePair;
import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.read.common.block.column.TimeColumn;
import org.apache.tsfile.read.filter.basic.Filter;
import org.apache.tsfile.read.filter.operator.TimeFilterOperators.TimeGt;
import org.apache.tsfile.read.filter.operator.TimeFilterOperators.TimeGtEq;
@@ -2218,6 +2222,49 @@ public Operator visitSort(SortNode node, LocalExecutionPlanContext context) {
getComparator(sortItemList, sortItemIndexList, sortItemDataTypeList));
}
+ @Override
+ public Operator visitInference(InferenceNode node, LocalExecutionPlanContext context) {
+ Operator child = node.getChild().accept(this, context);
+ OperatorContext operatorContext =
+ context
+ .getDriverContext()
+ .addOperatorContext(
+ context.getNextOperatorId(),
+ node.getPlanNodeId(),
+ org.apache.iotdb.db.queryengine.execution.operator.process.AI.InferenceOperator
+ .class
+ .getSimpleName());
+
+ ModelInferenceDescriptor modelInferenceDescriptor = node.getModelInferenceDescriptor();
+ ModelInformation modelInformation = modelInferenceDescriptor.getModelInformation();
+ int[] inputShape = modelInformation.getInputShape();
+ int[] outputShape = modelInformation.getOutputShape();
+ TSDataType[] inputTypes = modelInferenceDescriptor.getModelInformation().getInputDataType();
+ TSDataType[] outputTypes = modelInferenceDescriptor.getModelInformation().getOutputDataType();
+
+ long maxRetainedSize =
+ calculateSize(inputShape[0], inputTypes) + TimeColumn.SIZE_IN_BYTES_PER_POSITION;
+ long maxReturnSize = calculateSize(outputShape[0], outputTypes);
+
+ return new org.apache.iotdb.db.queryengine.execution.operator.process.AI.InferenceOperator(
+ operatorContext,
+ child,
+ modelInferenceDescriptor,
+ FragmentInstanceManager.getInstance().getModelInferenceExecutor(),
+ node.getInputColumnNames(),
+ node.getChild().getOutputColumnNames(),
+ maxRetainedSize,
+ maxReturnSize);
+ }
+
+ private long calculateSize(long rowNumber, TSDataType[] dataTypes) {
+ long size = 0;
+ for (int i = 0; i < dataTypes.length; i++) {
+ size += getOutputColumnSizePerLine(dataTypes[i]);
+ }
+ return size * rowNumber;
+ }
+
@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/PlanNodeType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java
index 061f22f171d5..b3e1d212041a 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
@@ -62,6 +62,7 @@
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedNonWritePlanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedWritePlanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeOperateSchemaQueueNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AI.InferenceNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ActiveRegionScanMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationMergeSortNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationNode;
@@ -197,9 +198,7 @@ public enum PlanNodeType {
LOGICAL_VIEW_SCHEMA_SCAN((short) 77),
ALTER_LOGICAL_VIEW((short) 78),
PIPE_ENRICHED_INSERT_DATA((short) 79),
-
- // NodeId 80 is used by IoTDB-ML which shouldn't be used.
-
+ INFERENCE((short) 80),
LAST_QUERY_TRANSFORM((short) 81),
TOP_K((short) 82),
COLUMN_INJECT((short) 83),
@@ -437,6 +436,8 @@ public static PlanNode deserialize(ByteBuffer buffer, short nodeType) {
return AlterLogicalViewNode.deserialize(buffer);
case 79:
return PipeEnrichedInsertNode.deserialize(buffer);
+ case 80:
+ return InferenceNode.deserialize(buffer);
case 81:
return LastQueryTransformNode.deserialize(buffer);
case 82:
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 2654487da54c..92261d7ad479 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
@@ -59,6 +59,7 @@
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedNonWritePlanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedWritePlanNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeOperateSchemaQueueNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AI.InferenceNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ActiveRegionScanMergeNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationMergeSortNode;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AggregationNode;
@@ -236,6 +237,10 @@ public R visitSingleDeviceView(SingleDeviceViewNode node, C context) {
return visitSingleChildProcess(node, context);
}
+ public R visitInference(InferenceNode node, C context) {
+ return visitSingleChildProcess(node, context);
+ }
+
public R visitExplainAnalyze(ExplainAnalyzeNode node, C context) {
return visitSingleChildProcess(node, context);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/AI/InferenceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/AI/InferenceNode.java
new file mode 100644
index 000000000000..95fe3437e788
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/AI/InferenceNode.java
@@ -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.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.AI;
+
+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.ModelInferenceDescriptor;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Objects;
+
+public class InferenceNode extends SingleChildProcessNode {
+
+ private final ModelInferenceDescriptor modelInferenceDescriptor;
+
+ // the column order in select item which reflects the real input order
+ private final List targetColumnNames;
+
+ public InferenceNode(
+ PlanNodeId id,
+ PlanNode child,
+ ModelInferenceDescriptor modelInferenceDescriptor,
+ List targetColumnNames) {
+ super(id, child);
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ this.targetColumnNames = targetColumnNames;
+ }
+
+ public InferenceNode(
+ PlanNodeId id,
+ ModelInferenceDescriptor modelInferenceDescriptor,
+ List inputColumnNames) {
+ super(id);
+ this.modelInferenceDescriptor = modelInferenceDescriptor;
+ this.targetColumnNames = inputColumnNames;
+ }
+
+ public ModelInferenceDescriptor getModelInferenceDescriptor() {
+ return modelInferenceDescriptor;
+ }
+
+ public List getInputColumnNames() {
+ return targetColumnNames;
+ }
+
+ @Override
+ public R accept(PlanVisitor visitor, C context) {
+ return visitor.visitInference(this, context);
+ }
+
+ @Override
+ public PlanNode clone() {
+ return new InferenceNode(getPlanNodeId(), child, modelInferenceDescriptor, targetColumnNames);
+ }
+
+ @Override
+ public List getOutputColumnNames() {
+ return modelInferenceDescriptor.getOutputColumnNames();
+ }
+
+ @Override
+ protected void serializeAttributes(ByteBuffer byteBuffer) {
+ PlanNodeType.INFERENCE.serialize(byteBuffer);
+ modelInferenceDescriptor.serialize(byteBuffer);
+ ReadWriteIOUtils.writeStringList(targetColumnNames, byteBuffer);
+ }
+
+ @Override
+ protected void serializeAttributes(DataOutputStream stream) throws IOException {
+ PlanNodeType.INFERENCE.serialize(stream);
+ modelInferenceDescriptor.serialize(stream);
+ ReadWriteIOUtils.writeStringList(targetColumnNames, stream);
+ }
+
+ public static InferenceNode deserialize(ByteBuffer buffer) {
+ ModelInferenceDescriptor modelInferenceDescriptor =
+ ModelInferenceDescriptor.deserialize(buffer);
+ List inputColumnNames = ReadWriteIOUtils.readStringList(buffer);
+ PlanNodeId planNodeId = PlanNodeId.deserialize(buffer);
+ return new InferenceNode(planNodeId, modelInferenceDescriptor, inputColumnNames);
+ }
+
+ @Override
+ public String toString() {
+ return "InferenceNode-" + 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;
+ }
+ InferenceNode that = (InferenceNode) o;
+ return modelInferenceDescriptor.equals(that.modelInferenceDescriptor)
+ && targetColumnNames.equals(that.targetColumnNames);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), modelInferenceDescriptor, targetColumnNames);
+ }
+}
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 000000000000..bf5f391d9e47
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/parameter/model/ModelInferenceDescriptor.java
@@ -0,0 +1,204 @@
+/*
+ * 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.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.model.ModelInformation;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindowParameter;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public class ModelInferenceDescriptor {
+
+ private final TEndPoint targetAINode;
+ private final ModelInformation modelInformation;
+ private List outputColumnNames;
+ private InferenceWindowParameter inferenceWindowParameter;
+ private Map inferenceAttributes;
+
+ public ModelInferenceDescriptor(TEndPoint targetAINode, ModelInformation modelInformation) {
+ this.targetAINode = targetAINode;
+ this.modelInformation = modelInformation;
+ }
+
+ private ModelInferenceDescriptor(ByteBuffer buffer) {
+ this.targetAINode =
+ new TEndPoint(ReadWriteIOUtils.readString(buffer), ReadWriteIOUtils.readInt(buffer));
+ this.modelInformation = ModelInformation.deserialize(buffer);
+ int outputColumnNamesSize = ReadWriteIOUtils.readInt(buffer);
+ if (outputColumnNamesSize == 0) {
+ this.outputColumnNames = null;
+ } else {
+ this.outputColumnNames = new ArrayList<>();
+ for (int i = 0; i < outputColumnNamesSize; i++) {
+ this.outputColumnNames.add(ReadWriteIOUtils.readString(buffer));
+ }
+ }
+ boolean hasInferenceWindowParameter = ReadWriteIOUtils.readBool(buffer);
+ if (hasInferenceWindowParameter) {
+ this.inferenceWindowParameter = InferenceWindowParameter.deserialize(buffer);
+ } else {
+ this.inferenceWindowParameter = null;
+ }
+ int inferenceAttributesSize = ReadWriteIOUtils.readInt(buffer);
+ if (inferenceAttributesSize == 0) {
+ this.inferenceAttributes = null;
+ } else {
+ this.inferenceAttributes = new HashMap<>();
+ for (int i = 0; i < inferenceAttributesSize; i++) {
+ this.inferenceAttributes.put(
+ ReadWriteIOUtils.readString(buffer), ReadWriteIOUtils.readString(buffer));
+ }
+ }
+ }
+
+ public void setInferenceAttributes(Map inferenceAttributes) {
+ this.inferenceAttributes = inferenceAttributes;
+ }
+
+ public Map getInferenceAttributes() {
+ return inferenceAttributes;
+ }
+
+ public void setInferenceWindowParameter(InferenceWindowParameter inferenceWindowParameter) {
+ this.inferenceWindowParameter = inferenceWindowParameter;
+ }
+
+ public InferenceWindowParameter getInferenceWindowParameter() {
+ return inferenceWindowParameter;
+ }
+
+ public ModelInformation getModelInformation() {
+ return modelInformation;
+ }
+
+ public TEndPoint getTargetAINode() {
+ return targetAINode;
+ }
+
+ public String getModelName() {
+ return modelInformation.getModelName();
+ }
+
+ public void setOutputColumnNames(List outputColumnNames) {
+ this.outputColumnNames = outputColumnNames;
+ }
+
+ public List getOutputColumnNames() {
+ return outputColumnNames;
+ }
+
+ public void serialize(ByteBuffer byteBuffer) {
+ ReadWriteIOUtils.write(targetAINode.ip, byteBuffer);
+ ReadWriteIOUtils.write(targetAINode.port, byteBuffer);
+ modelInformation.serialize(byteBuffer);
+ if (outputColumnNames == null) {
+ ReadWriteIOUtils.write(0, byteBuffer);
+ } else {
+ ReadWriteIOUtils.write(outputColumnNames.size(), byteBuffer);
+ for (String outputColumnName : outputColumnNames) {
+ ReadWriteIOUtils.write(outputColumnName, byteBuffer);
+ }
+ }
+ if (inferenceWindowParameter == null) {
+ ReadWriteIOUtils.write(false, byteBuffer);
+ } else {
+ ReadWriteIOUtils.write(true, byteBuffer);
+ inferenceWindowParameter.serialize(byteBuffer);
+ }
+ if (inferenceAttributes == null) {
+ ReadWriteIOUtils.write(0, byteBuffer);
+ } else {
+ ReadWriteIOUtils.write(inferenceAttributes.size(), byteBuffer);
+ for (Map.Entry entry : inferenceAttributes.entrySet()) {
+ ReadWriteIOUtils.write(entry.getKey(), byteBuffer);
+ ReadWriteIOUtils.write(entry.getValue(), byteBuffer);
+ }
+ }
+ }
+
+ public void serialize(DataOutputStream stream) throws IOException {
+ ReadWriteIOUtils.write(targetAINode.ip, stream);
+ ReadWriteIOUtils.write(targetAINode.port, stream);
+ modelInformation.serialize(stream);
+ if (outputColumnNames == null) {
+ ReadWriteIOUtils.write(0, stream);
+ } else {
+ ReadWriteIOUtils.write(outputColumnNames.size(), stream);
+ for (String outputColumnName : outputColumnNames) {
+ ReadWriteIOUtils.write(outputColumnName, stream);
+ }
+ }
+ if (inferenceWindowParameter == null) {
+ ReadWriteIOUtils.write(false, stream);
+ } else {
+ ReadWriteIOUtils.write(true, stream);
+ inferenceWindowParameter.serialize(stream);
+ }
+ if (inferenceAttributes == null) {
+ ReadWriteIOUtils.write(0, stream);
+ } else {
+ ReadWriteIOUtils.write(inferenceAttributes.size(), stream);
+ for (Map.Entry entry : inferenceAttributes.entrySet()) {
+ ReadWriteIOUtils.write(entry.getKey(), stream);
+ ReadWriteIOUtils.write(entry.getValue(), stream);
+ }
+ }
+ }
+
+ public static ModelInferenceDescriptor deserialize(ByteBuffer buffer) {
+ return new ModelInferenceDescriptor(buffer);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ ModelInferenceDescriptor that = (ModelInferenceDescriptor) o;
+ return targetAINode.equals(that.targetAINode)
+ && modelInformation.equals(that.modelInformation)
+ && outputColumnNames.equals(that.outputColumnNames)
+ && inferenceWindowParameter.equals(that.inferenceWindowParameter)
+ && inferenceAttributes.equals(that.inferenceAttributes);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ targetAINode,
+ modelInformation,
+ outputColumnNames,
+ inferenceWindowParameter,
+ inferenceAttributes);
+ }
+}
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 e7b7e3ffa663..d15677de60b7 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
@@ -74,6 +74,10 @@
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowTriggersStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowVariablesStatement;
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.ShowAINodesStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -283,6 +287,19 @@ public R visitAlterLogicalView(AlterLogicalViewStatement alterLogicalViewStateme
return visitStatement(alterLogicalViewStatement, context);
}
+ // AI Model
+ public R visitCreateModel(CreateModelStatement createModelStatement, C context) {
+ return visitStatement(createModelStatement, context);
+ }
+
+ public R visitDropModel(DropModelStatement dropModelStatement, C context) {
+ return visitStatement(dropModelStatement, context);
+ }
+
+ public R visitShowModels(ShowModelsStatement showModelsModelStatement, C context) {
+ return visitStatement(showModelsModelStatement, context);
+ }
+
/** Data Manipulation Language (DML) */
// Select Statement
@@ -445,6 +462,10 @@ public R visitShowConfigNodes(ShowConfigNodesStatement showConfigNodesStatement,
return visitStatement(showConfigNodesStatement, context);
}
+ public R visitShowAINodes(ShowAINodesStatement showAINodesStatement, C context) {
+ return visitStatement(showAINodesStatement, context);
+ }
+
public R visitShowVersion(ShowVersionStatement showVersionStatement, C context) {
return visitStatement(showVersionStatement, context);
}
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 e09bdc004150..2ff3e8149c89 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
@@ -26,6 +26,7 @@
import org.apache.iotdb.db.auth.AuthorityChecker;
import org.apache.iotdb.db.exception.sql.SemanticException;
import org.apache.iotdb.db.queryengine.execution.operator.window.WindowType;
+import org.apache.iotdb.db.queryengine.execution.operator.window.ainode.InferenceWindow;
import org.apache.iotdb.db.queryengine.plan.analyze.ExpressionAnalyzer;
import org.apache.iotdb.db.queryengine.plan.expression.Expression;
import org.apache.iotdb.db.queryengine.plan.expression.leaf.TimeSeriesOperand;
@@ -53,8 +54,10 @@
import java.util.ArrayList;
import java.util.Collections;
+import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
+import java.util.Map;
import java.util.Set;
import static org.apache.iotdb.db.utils.constant.SqlConstant.COUNT_TIME;
@@ -133,6 +136,57 @@ public class QueryStatement extends AuthorityInformationStatement {
// we can skip the query
private boolean isResultSetEmpty = false;
+ // [IoTDB-AI] used for model inference, which will be removed in the future
+ private String modelName;
+ private boolean hasModelInference = false;
+ private InferenceWindow inferenceWindow = null;
+ private Map inferenceAttribute = null;
+
+ public void setModelName(String modelName) {
+ this.modelName = modelName;
+ }
+
+ public String getModelName() {
+ return modelName;
+ }
+
+ public void setHasModelInference(boolean hasModelInference) {
+ this.hasModelInference = hasModelInference;
+ }
+
+ public boolean hasModelInference() {
+ return hasModelInference;
+ }
+
+ public void setInferenceWindow(InferenceWindow inferenceWindow) {
+ this.inferenceWindow = inferenceWindow;
+ }
+
+ public boolean isSetInferenceWindow() {
+ return this.inferenceWindow != null;
+ }
+
+ public InferenceWindow getInferenceWindow() {
+ return inferenceWindow;
+ }
+
+ public void addInferenceAttribute(String key, String value) {
+ if (inferenceAttribute == null) {
+ inferenceAttribute = new HashMap<>();
+ }
+ inferenceAttribute.put(key, value);
+ }
+
+ public Map getInferenceAttributes() {
+ return inferenceAttribute;
+ }
+
+ public boolean hasInferenceAttributes() {
+ return inferenceAttribute != null;
+ }
+
+ // [IoTDB-AI] END
+
public QueryStatement() {
this.statementType = StatementType.QUERY;
}
@@ -546,6 +600,16 @@ public void setLastLevelUseWildcard(boolean lastLevelUseWildcard) {
@SuppressWarnings({"squid:S3776", "squid:S6541"}) // Suppress high Cognitive Complexity warning
public void semanticCheck() {
+
+ if (hasModelInference) {
+ if (isAlignByDevice()) {
+ throw new SemanticException("Model inference does not support align by device now.");
+ }
+ if (isSelectInto()) {
+ throw new SemanticException("Model inference does not support select into now.");
+ }
+ }
+
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
new file mode 100644
index 000000000000..2f43a1cfe762
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/CreateModelStatement.java
@@ -0,0 +1,85 @@
+/*
+ * 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.statement.metadata.model;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.auth.entity.PrivilegeType;
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.db.auth.AuthorityChecker;
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
+import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.Statement;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import java.util.Collections;
+import java.util.List;
+
+public class CreateModelStatement extends Statement implements IConfigStatement {
+
+ private String modelName;
+
+ private String uri;
+
+ public CreateModelStatement() {
+ // do nothing
+ }
+
+ public String getUri() {
+ return uri;
+ }
+
+ public String getModelName() {
+ return modelName;
+ }
+
+ public void setUri(String uri) {
+ this.uri = uri;
+ }
+
+ public void setModelName(String modelName) {
+ this.modelName = modelName;
+ }
+
+ @Override
+ public List getPaths() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public QueryType getQueryType() {
+ return QueryType.WRITE;
+ }
+
+ @Override
+ public TSStatus checkPermissionBeforeProcess(String userName) {
+ if (AuthorityChecker.SUPER_USER.equals(userName)) {
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+ return AuthorityChecker.getTSStatus(
+ AuthorityChecker.checkSystemPermission(userName, PrivilegeType.USE_MODEL.ordinal()),
+ PrivilegeType.USE_MODEL);
+ }
+
+ @Override
+ public R accept(StatementVisitor visitor, C context) {
+ return visitor.visitCreateModel(this, context);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/DropModelStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/DropModelStatement.java
new file mode 100644
index 000000000000..3e207608241d
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/DropModelStatement.java
@@ -0,0 +1,71 @@
+/*
+ * 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.statement.metadata.model;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.auth.entity.PrivilegeType;
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.db.auth.AuthorityChecker;
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
+import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.Statement;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import java.util.Collections;
+import java.util.List;
+
+public class DropModelStatement extends Statement implements IConfigStatement {
+
+ private final String modelName;
+
+ public DropModelStatement(String modelName) {
+ this.modelName = modelName;
+ }
+
+ public String getModelName() {
+ return modelName;
+ }
+
+ @Override
+ public List getPaths() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public QueryType getQueryType() {
+ return QueryType.WRITE;
+ }
+
+ @Override
+ public TSStatus checkPermissionBeforeProcess(String userName) {
+ if (AuthorityChecker.SUPER_USER.equals(userName)) {
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+ return AuthorityChecker.getTSStatus(
+ AuthorityChecker.checkSystemPermission(userName, PrivilegeType.USE_MODEL.ordinal()),
+ PrivilegeType.USE_MODEL);
+ }
+
+ @Override
+ public R accept(StatementVisitor visitor, C context) {
+ return visitor.visitDropModel(this, context);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowAINodesStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowAINodesStatement.java
new file mode 100644
index 000000000000..602d0e01465b
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowAINodesStatement.java
@@ -0,0 +1,40 @@
+/*
+ * 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.statement.metadata.model;
+
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
+import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
+import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowStatement;
+
+public class ShowAINodesStatement extends ShowStatement implements IConfigStatement {
+
+ public ShowAINodesStatement() {}
+
+ @Override
+ public QueryType getQueryType() {
+ return QueryType.READ;
+ }
+
+ @Override
+ public R accept(StatementVisitor visitor, C context) {
+ return visitor.visitShowAINodes(this, context);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowModelsStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowModelsStatement.java
new file mode 100644
index 000000000000..0b810b49ad82
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/metadata/model/ShowModelsStatement.java
@@ -0,0 +1,79 @@
+/*
+ * 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.statement.metadata.model;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.auth.entity.PrivilegeType;
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.db.auth.AuthorityChecker;
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
+import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.Statement;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import java.util.Collections;
+import java.util.List;
+
+public class ShowModelsStatement extends Statement implements IConfigStatement {
+
+ private String modelName;
+
+ public ShowModelsStatement() {
+ // do nothing
+ }
+
+ public void setModelName(String modelName) {
+ this.modelName = modelName;
+ }
+
+ public boolean isSetModelName() {
+ return modelName != null;
+ }
+
+ public String getModelName() {
+ return modelName;
+ }
+
+ @Override
+ public List getPaths() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public QueryType getQueryType() {
+ return QueryType.READ;
+ }
+
+ @Override
+ public TSStatus checkPermissionBeforeProcess(String userName) {
+ if (AuthorityChecker.SUPER_USER.equals(userName)) {
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+ return AuthorityChecker.getTSStatus(
+ AuthorityChecker.checkSystemPermission(userName, PrivilegeType.USE_MODEL.ordinal()),
+ PrivilegeType.USE_MODEL);
+ }
+
+ @Override
+ public R accept(StatementVisitor visitor, C context) {
+ return visitor.visitShowModels(this, context);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCService.java
new file mode 100644
index 000000000000..5ec49e756602
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCService.java
@@ -0,0 +1,94 @@
+/*
+ * 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.service;
+
+import org.apache.iotdb.commons.concurrent.ThreadName;
+import org.apache.iotdb.commons.exception.runtime.RPCServiceException;
+import org.apache.iotdb.commons.service.ServiceType;
+import org.apache.iotdb.commons.service.ThriftService;
+import org.apache.iotdb.commons.service.ThriftServiceThread;
+import org.apache.iotdb.db.conf.IoTDBConfig;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.protocol.thrift.handler.AINodeRPCServiceThriftHandler;
+import org.apache.iotdb.db.protocol.thrift.impl.AINodeRPCServiceImpl;
+import org.apache.iotdb.mpp.rpc.thrift.IAINodeInternalRPCService;
+import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory;
+
+public class AINodeRPCService extends ThriftService implements AINodeRPCServiceMBean {
+
+ private AINodeRPCServiceImpl impl;
+
+ private AINodeRPCService() {}
+
+ @Override
+ public ServiceType getID() {
+ return ServiceType.AINode_RPC_SERVICE;
+ }
+
+ @Override
+ public void initTProcessor() {
+ impl = new AINodeRPCServiceImpl();
+ initSyncedServiceImpl(null);
+ processor = new IAINodeInternalRPCService.Processor<>(impl);
+ }
+
+ @Override
+ public void initThriftServiceThread()
+ throws IllegalAccessException, InstantiationException, ClassNotFoundException {
+ try {
+ IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+ thriftServiceThread =
+ new ThriftServiceThread(
+ processor,
+ getID().getName(),
+ ThreadName.AINODE_RPC_SERVICE.getName(),
+ getBindIP(),
+ getBindPort(),
+ config.getRpcMaxConcurrentClientNum(),
+ config.getThriftServerAwaitTimeForStopService(),
+ new AINodeRPCServiceThriftHandler(impl),
+ config.isRpcThriftCompressionEnable(),
+ DeepCopyRpcTransportFactory.INSTANCE);
+ } catch (RPCServiceException e) {
+ throw new IllegalAccessException(e.getMessage());
+ }
+ thriftServiceThread.setName(ThreadName.AINODE_RPC_SERVICE.getName());
+ }
+
+ @Override
+ public String getBindIP() {
+ return IoTDBDescriptor.getInstance().getConfig().getRpcAddress();
+ }
+
+ @Override
+ public int getBindPort() {
+ return IoTDBDescriptor.getInstance().getConfig().getAINodePort();
+ }
+
+ private static class AINodeRPCServiceHolder {
+ private static final AINodeRPCService INSTANCE = new AINodeRPCService();
+
+ private AINodeRPCServiceHolder() {}
+ }
+
+ public static AINodeRPCService getInstance() {
+ return AINodeRPCServiceHolder.INSTANCE;
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCServiceMBean.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCServiceMBean.java
new file mode 100644
index 000000000000..f4f51c0caa27
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/AINodeRPCServiceMBean.java
@@ -0,0 +1,22 @@
+/*
+ * 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.service;
+
+public interface AINodeRPCServiceMBean {}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java
index d1ea85fe4c71..89bd0df7c286 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java
@@ -684,6 +684,10 @@ private void setUp() throws StartupException {
private void setUpRPCService() throws StartupException {
// Start InternalRPCService to indicate that the current DataNode can accept cluster scheduling
registerManager.register(DataNodeInternalRPCService.getInstance());
+ // Start InternalRPCService to indicate that the current DataNode can accept request from AINode
+ if (config.isEnableAINodeService()) {
+ registerManager.register(AINodeRPCService.getInstance());
+ }
// Notice: During the period between starting the internal RPC service
// and starting the client RPC service , some requests may fail because
diff --git a/iotdb-core/node-commons/pom.xml b/iotdb-core/node-commons/pom.xml
index 56ab07bd99d7..262abd5e2039 100644
--- a/iotdb-core/node-commons/pom.xml
+++ b/iotdb-core/node-commons/pom.xml
@@ -182,6 +182,12 @@
org.apache.ratis
ratis-common
+