AINode: Fix built-model inference error (apache#13985)

Caideyipi · Nov 4, 2024 · be720d2 · be720d2
1 parent 005295b
commit be720d2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
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
@@ -179,6 +179,8 @@ public void ModelOperationTest() {
   public void callInferenceTest() {
     String sql = "CALL INFERENCE(identity, \"select s0,s1,s2 from root.AI.data\")";
     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)";
     try (Connection connection = EnvFactory.getEnv().getConnection();
         Statement statement = connection.createStatement()) {
 
@@ -215,6 +217,17 @@ public void callInferenceTest() {
         }
         assertEquals(7, count);
       }
+
+      try (ResultSet resultSet = statement.executeQuery(sql3)) {
+        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+        checkHeader(resultSetMetaData, "output0,output1,output2");
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(3, count);
+      }
+
     } catch (SQLException e) {
       fail(e.getMessage());
     }

diff --git a/iotdb-core/ainode/iotdb/ainode/model/built_in_model_factory.py b/iotdb-core/ainode/iotdb/ainode/model/built_in_model_factory.py
@@ -99,7 +99,7 @@ def fetch_built_in_model(model_id, inference_attributes):
     else:
         raise BuiltInModelNotSupportError(model_id)
 
-    return model, attributes
+    return model
 
 
 class Attribute(object):