[AINode] Support output time column for model inference

ycycse · web-flow · commit f16252d8105c · 2025-03-13T09:19:45.000+08:00
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
@@ -178,21 +178,22 @@ public void ModelOperationTest() {
 
   @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 void callInferenceTest() {
 
       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
@@ -36,6 +36,7 @@
 
 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,13 +82,20 @@ 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,
       ModelInferenceDescriptor modelInferenceDescriptor,
       ExecutorService modelInferenceExecutor,
       List<String> targetColumnNames,
       List<String> inputColumnNames,
+      boolean generateTimeColumn,
       long maxRetainedSize,
       long maxReturnSize) {
     this.operatorContext = operatorContext;
@@ -106,6 +114,14 @@ public InferenceOperator(
     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 boolean hasNext() throws Exception {
     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 TsBlock next() throws Exception {
 
       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 TsBlock next() throws Exception {
 
         finished = true;
         TsBlock resultTsBlock = serde.deserialize(inferenceResp.inferenceResult.get(0));
+        if (generateTimeColumn) {
+          fillTimeColumn(resultTsBlock);
+        }
         results = inferenceResp.inferenceResult;
         resultIndex++;
         return resultTsBlock;
@@ -194,7 +225,12 @@ private void appendTsBlockToBuilder(TsBlock inputTsBlock) {
     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 @@ private TsBlock preProcess(TsBlock inputTsBlock) {
 
   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
@@ -1695,7 +1695,8 @@ static void analyzeOutput(
           .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
@@ -4527,6 +4527,9 @@ public Statement visitCallInference(IoTDBSqlParser.CallInferenceContext ctx) {
                 "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
@@ -1383,6 +1383,7 @@ public LogicalPlanBuilder planInference(Analysis analysis) {
             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
@@ -234,7 +234,7 @@ public PlanNode visitQuery(QueryStatement queryStatement, MPPQueryContext contex
     }
 
     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
@@ -2310,6 +2310,7 @@ public Operator visitInference(InferenceNode node, LocalExecutionPlanContext con
         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
@@ -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,14 +73,19 @@ public List<String> getInputColumnNames() {
     return targetColumnNames;
   }
 
+  public boolean isGenerateTimeColumn() {
+    return generateTimeColumn;
+  }
+
   @Override
   public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
     return visitor.visitInference(this, context);
   }
 
   @Override
   public PlanNode clone() {
-    return new InferenceNode(getPlanNodeId(), child, modelInferenceDescriptor, targetColumnNames);
+    return new InferenceNode(
+        getPlanNodeId(), child, modelInferenceDescriptor, generateTimeColumn, targetColumnNames);
   }
 
   @Override
@@ -87,22 +97,26 @@ public List<String> getOutputColumnNames() {
   protected void serializeAttributes(ByteBuffer byteBuffer) {
     PlanNodeType.INFERENCE.serialize(byteBuffer);
     modelInferenceDescriptor.serialize(byteBuffer);
+    ReadWriteIOUtils.write(generateTimeColumn, byteBuffer);
     ReadWriteIOUtils.writeStringList(targetColumnNames, byteBuffer);
   }
 
   @Override
   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
@@ -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;
   }

Original file line number	Diff line number	Diff line change
`@@ -1695,7 +1695,8 @@ static void analyzeOutput(`
`1695`	`1695`	`.getModelInferenceDescriptor()`
`1696`	`1696`	`.setOutputColumnNames(`
`1697`	`1697`	`columnHeaders.stream().map(ColumnHeader::getColumnName).collect(Collectors.toList()));`
`1698`		`- analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, true));`
	`1698`	`+ boolean isIgnoreTimestamp = !queryStatement.isGenerateTime();`
	`1699`	`+ analysis.setRespDatasetHeader(new DatasetHeader(columnHeaders, isIgnoreTimestamp));`
`1699`	`1700`	`return;`
`1700`	`1701`	`}`
`1701`	`1702`
Original file line number	Diff line number	Diff line change
`@@ -234,7 +234,7 @@ public PlanNode visitQuery(QueryStatement queryStatement, MPPQueryContext contex`
`234`	`234`	`}`
`235`	`235`
`236`	`236`	`if (queryStatement.hasModelInference()) {`
`237`		`- planBuilder.planInference(analysis);`
	`237`	`+ planBuilder = planBuilder.planInference(analysis);`
`238`	`238`	`}`
`239`	`239`
`240`	`240`	`// plan select into`
Original file line number	Diff line number	Diff line change
`@@ -2310,6 +2310,7 @@ public Operator visitInference(InferenceNode node, LocalExecutionPlanContext con`
`2310`	`2310`	`FragmentInstanceManager.getInstance().getModelInferenceExecutor(),`
`2311`	`2311`	`node.getInputColumnNames(),`
`2312`	`2312`	`node.getChild().getOutputColumnNames(),`
	`2313`	`+ node.isGenerateTimeColumn(),`
`2313`	`2314`	`maxRetainedSize,`
`2314`	`2315`	`maxReturnSize);`
`2315`	`2316`	`}`