This is an automated email from the ASF dual-hosted git repository. ycycse pushed a commit to branch appendTime in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 57a15129bf0cf8cbb00a6535ce4803ec7b7b9195 Author: YangCaiyin <[email protected]> AuthorDate: Thu Mar 13 09:19:45 2025 +0800 [AINode] Support output time column for model inference (cherry picked from commit f16252d8105cab29bbe28b5d8ddf3d3afda6d85b) --- .../org/apache/iotdb/ainode/it/AINodeBasicIT.java | 15 ++++---- .../operator/process/ai/InferenceOperator.java | 42 +++++++++++++++++++++- .../queryengine/plan/analyze/AnalyzeVisitor.java | 3 +- .../db/queryengine/plan/parser/ASTVisitor.java | 3 ++ .../plan/planner/LogicalPlanBuilder.java | 1 + .../plan/planner/LogicalPlanVisitor.java | 2 +- .../plan/planner/OperatorTreeGenerator.java | 1 + .../plan/node/process/AI/InferenceNode.java | 18 ++++++++-- .../plan/statement/crud/QueryStatement.java | 9 +++++ 9 files changed, 82 insertions(+), 12 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java index 3b54b5c4409..07d29c0d224 100644 --- a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java @@ -178,21 +178,22 @@ public class AINodeBasicIT { @Test public void callInferenceTest() { - String sql = "CALL INFERENCE(identity, \"select s0,s1,s2 from root.AI.data\")"; + String sql = + "CALL INFERENCE(identity, \"select s0,s1,s2 from root.AI.data\", generateTime=true)"; String sql2 = "CALL INFERENCE(identity, \"select s2,s0,s1 from root.AI.data\")"; String sql3 = - "CALL INFERENCE(_NaiveForecaster, \"select s0 from root.AI.data\", predict_length=3)"; + "CALL INFERENCE(_NaiveForecaster, \"select s0 from root.AI.data\", predict_length=3, generateTime=true)"; try (Connection connection = EnvFactory.getEnv().getConnection(); Statement statement = connection.createStatement()) { try (ResultSet resultSet = statement.executeQuery(sql)) { ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); - checkHeader(resultSetMetaData, "output0,output1,output2"); + checkHeader(resultSetMetaData, "Time,output0,output1,output2"); int count = 0; while (resultSet.next()) { - float s0 = resultSet.getFloat(1); - float s1 = resultSet.getFloat(2); - float s2 = resultSet.getFloat(3); + float s0 = resultSet.getFloat(2); + float s1 = resultSet.getFloat(3); + float s2 = resultSet.getFloat(4); assertEquals(s0, count + 1.0, 0.0001); assertEquals(s1, count + 2.0, 0.0001); @@ -221,7 +222,7 @@ public class AINodeBasicIT { try (ResultSet resultSet = statement.executeQuery(sql3)) { ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); - checkHeader(resultSetMetaData, "output0,output1,output2"); + checkHeader(resultSetMetaData, "Time,output0,output1,output2"); int count = 0; while (resultSet.next()) { count++; 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 index 08cae2dc5b1..9bdb57dc5bd 100644 --- 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 @@ -36,6 +36,7 @@ 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.Column; import org.apache.tsfile.block.column.ColumnBuilder; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.read.common.block.TsBlock; @@ -81,6 +82,12 @@ public class InferenceOperator implements ProcessOperator { private final TsBlockSerde serde = new TsBlockSerde(); private InferenceWindowType windowType = null; + private final boolean generateTimeColumn; + private long maxTimestamp; + private long minTimestamp; + private long interval; + private long currentRowIndex; + public InferenceOperator( OperatorContext operatorContext, Operator child, @@ -88,6 +95,7 @@ public class InferenceOperator implements ProcessOperator { ExecutorService modelInferenceExecutor, List<String> targetColumnNames, List<String> inputColumnNames, + boolean generateTimeColumn, long maxRetainedSize, long maxReturnSize) { this.operatorContext = operatorContext; @@ -106,6 +114,14 @@ public class InferenceOperator implements ProcessOperator { if (modelInferenceDescriptor.getInferenceWindowParameter() != null) { windowType = modelInferenceDescriptor.getInferenceWindowParameter().getWindowType(); } + + if (generateTimeColumn) { + this.interval = 0; + this.minTimestamp = Long.MAX_VALUE; + this.maxTimestamp = Long.MIN_VALUE; + this.currentRowIndex = 0; + } + this.generateTimeColumn = generateTimeColumn; } @Override @@ -140,6 +156,15 @@ public class InferenceOperator implements ProcessOperator { return !finished || (results != null && results.size() != resultIndex); } + private void fillTimeColumn(TsBlock tsBlock) { + Column timeColumn = tsBlock.getTimeColumn(); + long[] time = timeColumn.getLongs(); + for (int i = 0; i < time.length; i++) { + time[i] = maxTimestamp + interval * currentRowIndex; + currentRowIndex++; + } + } + @Override public TsBlock next() throws Exception { if (inferenceExecutionFuture == null) { @@ -156,6 +181,9 @@ public class InferenceOperator implements ProcessOperator { if (results != null && resultIndex != results.size()) { TsBlock tsBlock = serde.deserialize(results.get(resultIndex)); + if (generateTimeColumn) { + fillTimeColumn(tsBlock); + } resultIndex++; return tsBlock; } @@ -177,6 +205,9 @@ public class InferenceOperator implements ProcessOperator { finished = true; TsBlock resultTsBlock = serde.deserialize(inferenceResp.inferenceResult.get(0)); + if (generateTimeColumn) { + fillTimeColumn(resultTsBlock); + } results = inferenceResp.inferenceResult; resultIndex++; return resultTsBlock; @@ -194,7 +225,12 @@ public class InferenceOperator implements ProcessOperator { ColumnBuilder[] columnBuilders = inputTsBlockBuilder.getValueColumnBuilders(); totalRow += inputTsBlock.getPositionCount(); for (int i = 0; i < inputTsBlock.getPositionCount(); i++) { - timeColumnBuilder.writeLong(inputTsBlock.getTimeByIndex(i)); + long timestamp = inputTsBlock.getTimeByIndex(i); + if (generateTimeColumn) { + minTimestamp = Math.min(minTimestamp, timestamp); + maxTimestamp = Math.max(maxTimestamp, timestamp); + } + timeColumnBuilder.writeLong(timestamp); for (int columnIndex = 0; columnIndex < inputTsBlock.getValueColumnCount(); columnIndex++) { columnBuilders[columnIndex].write(inputTsBlock.getColumn(columnIndex), i); } @@ -259,6 +295,10 @@ public class InferenceOperator implements ProcessOperator { private void submitInferenceTask() { + if (generateTimeColumn) { + interval = (maxTimestamp - minTimestamp) / totalRow; + } + TsBlock inputTsBlock = inputTsBlockBuilder.build(); TsBlock finalInputTsBlock = preProcess(inputTsBlock); 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 cd9fd65cfd9..a8b1d11b73a 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 @@ -1692,7 +1692,8 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext> .getModelInferenceDescriptor() .setOutputColumnNames( columnHeaders.stream().map(ColumnHeader::getColumnName).collect(Collectors.toList())); - analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, true)); + boolean isIgnoreTimestamp = !queryStatement.isGenerateTime(); + analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, isIgnoreTimestamp)); return; } 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 68a6b32ea09..e6b49403ab2 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 @@ -4477,6 +4477,9 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> { "Window Function(e.g. HEAD, TAIL, COUNT) should be set in value when key is 'WINDOW' in CALL INFERENCE"); } parseWindowFunctionInInference(valueContext.windowFunction(), statement); + } else if (paramKey.equalsIgnoreCase("GENERATETIME")) { + statement.setGenerateTime( + Boolean.parseBoolean(parseAttributeValue(valueContext.attributeValue()))); } else { statement.addInferenceAttribute( paramKey, parseAttributeValue(valueContext.attributeValue())); 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 3c234e8ba44..4c8eb7f1e39 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 @@ -1377,6 +1377,7 @@ public class LogicalPlanBuilder { context.getQueryId().genPlanNodeId(), root, analysis.getModelInferenceDescriptor(), + !analysis.getRespDatasetHeader().isIgnoreTimestamp(), analysis.getOutputExpressions().stream() .map(expressionStringPair -> expressionStringPair.left.getExpressionString()) .collect(Collectors.toList())); 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 86bc69f27e3..50f43f364fa 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 @@ -231,7 +231,7 @@ public class LogicalPlanVisitor extends StatementVisitor<PlanNode, MPPQueryConte } if (queryStatement.hasModelInference()) { - planBuilder.planInference(analysis); + planBuilder = planBuilder.planInference(analysis); } // plan select into 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 14f3803c1a1..c9405ea0bc6 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 @@ -2254,6 +2254,7 @@ public class OperatorTreeGenerator extends PlanVisitor<Operator, LocalExecutionP FragmentInstanceManager.getInstance().getModelInferenceExecutor(), node.getInputColumnNames(), node.getChild().getOutputColumnNames(), + node.isGenerateTimeColumn(), maxRetainedSize, maxReturnSize); } 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 index 95fe3437e78..09205c9eb56 100644 --- 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 @@ -40,24 +40,29 @@ public class InferenceNode extends SingleChildProcessNode { // the column order in select item which reflects the real input order private final List<String> targetColumnNames; + private boolean generateTimeColumn = false; public InferenceNode( PlanNodeId id, PlanNode child, ModelInferenceDescriptor modelInferenceDescriptor, + boolean generateTimeColumn, List<String> targetColumnNames) { super(id, child); this.modelInferenceDescriptor = modelInferenceDescriptor; this.targetColumnNames = targetColumnNames; + this.generateTimeColumn = generateTimeColumn; } public InferenceNode( PlanNodeId id, ModelInferenceDescriptor modelInferenceDescriptor, + boolean generateTimeColumn, List<String> inputColumnNames) { super(id); this.modelInferenceDescriptor = modelInferenceDescriptor; this.targetColumnNames = inputColumnNames; + this.generateTimeColumn = generateTimeColumn; } public ModelInferenceDescriptor getModelInferenceDescriptor() { @@ -68,6 +73,10 @@ public class InferenceNode extends SingleChildProcessNode { return targetColumnNames; } + public boolean isGenerateTimeColumn() { + return generateTimeColumn; + } + @Override public <R, C> R accept(PlanVisitor<R, C> visitor, C context) { return visitor.visitInference(this, context); @@ -75,7 +84,8 @@ public class InferenceNode extends SingleChildProcessNode { @Override public PlanNode clone() { - return new InferenceNode(getPlanNodeId(), child, modelInferenceDescriptor, targetColumnNames); + return new InferenceNode( + getPlanNodeId(), child, modelInferenceDescriptor, generateTimeColumn, targetColumnNames); } @Override @@ -87,6 +97,7 @@ public class InferenceNode extends SingleChildProcessNode { protected void serializeAttributes(ByteBuffer byteBuffer) { PlanNodeType.INFERENCE.serialize(byteBuffer); modelInferenceDescriptor.serialize(byteBuffer); + ReadWriteIOUtils.write(generateTimeColumn, byteBuffer); ReadWriteIOUtils.writeStringList(targetColumnNames, byteBuffer); } @@ -94,15 +105,18 @@ public class InferenceNode extends SingleChildProcessNode { protected void serializeAttributes(DataOutputStream stream) throws IOException { PlanNodeType.INFERENCE.serialize(stream); modelInferenceDescriptor.serialize(stream); + ReadWriteIOUtils.write(generateTimeColumn, stream); ReadWriteIOUtils.writeStringList(targetColumnNames, stream); } public static InferenceNode deserialize(ByteBuffer buffer) { ModelInferenceDescriptor modelInferenceDescriptor = ModelInferenceDescriptor.deserialize(buffer); + boolean generateTimeColumn = ReadWriteIOUtils.readBool(buffer); List<String> inputColumnNames = ReadWriteIOUtils.readStringList(buffer); PlanNodeId planNodeId = PlanNodeId.deserialize(buffer); - return new InferenceNode(planNodeId, modelInferenceDescriptor, inputColumnNames); + return new InferenceNode( + planNodeId, modelInferenceDescriptor, generateTimeColumn, inputColumnNames); } @Override 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 2ff3e8149c8..6d36d4fd64e 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 @@ -139,9 +139,18 @@ public class QueryStatement extends AuthorityInformationStatement { // [IoTDB-AI] used for model inference, which will be removed in the future private String modelName; private boolean hasModelInference = false; + private boolean generateTime = false; private InferenceWindow inferenceWindow = null; private Map<String, String> inferenceAttribute = null; + public void setGenerateTime(boolean generateTime) { + this.generateTime = generateTime; + } + + public boolean isGenerateTime() { + return generateTime; + } + public void setModelName(String modelName) { this.modelName = modelName; }
