Support built-in forecast function through UDTF for tree model (#15682)

(cherry picked from commit c158092)

Author: ycycse

integration-test/src/main/java/org/apache/iotdb/itbase/constant/BuiltinAggregationFunctionEnum.java

@@ -56,8 +56,11 @@ public String getFunctionName() {
 
   private static final Set<String> NATIVE_FUNCTION_NAMES =
       new HashSet<>(
-          Arrays.stream(org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction.values())
-              .map(org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction::getFunctionName)
+          Arrays.stream(
+                  org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction.values())
+              .map(
+                  org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction
+                      ::getFunctionName)
               .collect(Collectors.toList()));
 
   public static Set<String> getNativeFunctionNames() {

integration-test/src/main/java/org/apache/iotdb/itbase/constant/BuiltinTimeSeriesGeneratingFunctionEnum.java

@@ -74,6 +74,7 @@ public enum BuiltinTimeSeriesGeneratingFunctionEnum {
   EQUAL_SIZE_BUCKET_OUTLIER_SAMPLE("EQUAL_SIZE_BUCKET_OUTLIER_SAMPLE"),
   JEXL("JEXL"),
   MASTER_REPAIR("MASTER_REPAIR"),
+  FORECAST("FORECAST"),
   M4("M4");
 
   private final String functionName;

integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentInferenceIT.java

@@ -65,7 +65,7 @@ public static void tearDown() throws Exception {
 
   private static void prepareDataForTreeModel() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
-        Statement statement = connection.createStatement()) {
+         Statement statement = connection.createStatement()) {
       statement.execute("CREATE DATABASE root.AI");
       statement.execute("CREATE TIMESERIES root.AI.s WITH DATATYPE=DOUBLE, ENCODING=RLE");
       for (int i = 0; i < 2880; i++) {
@@ -79,7 +79,7 @@ private static void prepareDataForTreeModel() throws SQLException {
 
   private static void prepareDataForTableModel() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
-        Statement statement = connection.createStatement()) {
+         Statement statement = connection.createStatement()) {
       statement.execute("CREATE DATABASE root");
       statement.execute("CREATE TABLE root.AI (s DOUBLE FIELD)");
       for (int i = 0; i < 2880; i++) {
@@ -99,7 +99,7 @@ public void concurrentGPUCallInferenceTest() throws SQLException, InterruptedExc
   private void concurrentGPUCallInferenceTest(String modelId)
       throws SQLException, InterruptedException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
-        Statement statement = connection.createStatement()) {
+         Statement statement = connection.createStatement()) {
       final int threadCnt = 10;
       final int loop = 100;
       final int predictLength = 512;
@@ -118,15 +118,23 @@ private void concurrentGPUCallInferenceTest(String modelId)
     }
   }
 
+  String forecastTableFunctionSql =
+      "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time), predict_length=>%d";
+  String forecastUDTFSql =
+      "SELECT forecast(s, 'MODEL_ID'='%s', 'PREDICT_LENGTH'='%d') FROM root.AI";
+
   @Test
   public void concurrentGPUForecastTest() throws SQLException, InterruptedException {
-    concurrentGPUForecastTest("timer_xl");
-    concurrentGPUForecastTest("sundial");
+    concurrentGPUForecastTest("timer_xl", forecastUDTFSql);
+    concurrentGPUForecastTest("sundial", forecastUDTFSql);
+    concurrentGPUForecastTest("timer_xl", forecastTableFunctionSql);
+    concurrentGPUForecastTest("sundial", forecastTableFunctionSql);
   }
 
-  public void concurrentGPUForecastTest(String modelId) throws SQLException, InterruptedException {
+  public void concurrentGPUForecastTest(String modelId, String selectSql)
+      throws SQLException, InterruptedException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
-        Statement statement = connection.createStatement()) {
+         Statement statement = connection.createStatement()) {
       final int threadCnt = 10;
       final int loop = 100;
       final int predictLength = 512;
@@ -136,9 +144,7 @@ public void concurrentGPUForecastTest(String modelId) throws SQLException, Inter
       long startTime = System.currentTimeMillis();
       concurrentInference(
           statement,
-          String.format(
-              "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time), predict_length=>%d",
-              modelId, predictLength),
+          String.format(selectSql, modelId, predictLength),
           threadCnt,
           loop,
           predictLength);
@@ -158,7 +164,7 @@ private void checkModelOnSpecifiedDevice(Statement statement, String modelId, St
     for (int retry = 0; retry < 200; retry++) {
       Set<String> foundDevices = new HashSet<>();
       try (final ResultSet resultSet =
-          statement.executeQuery(String.format("SHOW LOADED MODELS '%s'", device))) {
+               statement.executeQuery(String.format("SHOW LOADED MODELS '%s'", device))) {
         while (resultSet.next()) {
           String deviceId = resultSet.getString("DeviceId");
           String loadedModelId = resultSet.getString("ModelId");

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationUtil.java

@@ -19,7 +19,6 @@
 
 package org.apache.iotdb.db.queryengine.execution.operator;
 
-import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
 import org.apache.iotdb.db.queryengine.execution.aggregation.TreeAggregator;
 import org.apache.iotdb.db.queryengine.execution.aggregation.timerangeiterator.ITimeRangeIterator;
 import org.apache.iotdb.db.queryengine.execution.aggregation.timerangeiterator.SingleTimeWindowIterator;
@@ -29,6 +28,7 @@
 import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.AggregationDescriptor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByTimeParameter;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction;
 import org.apache.iotdb.db.queryengine.statistics.StatisticsManager;
 
 import org.apache.tsfile.block.column.Column;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionAnalyzer.java

@@ -24,8 +24,6 @@
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.path.PathPatternTree;
 import org.apache.iotdb.commons.schema.column.ColumnHeader;
-import org.apache.iotdb.commons.udf.builtin.BuiltinScalarFunction;
-import org.apache.iotdb.commons.udf.builtin.BuiltinTimeSeriesGeneratingFunction;
 import org.apache.iotdb.db.exception.sql.SemanticException;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.schematree.ISchemaTree;
@@ -56,6 +54,8 @@
 import org.apache.iotdb.db.queryengine.plan.expression.visitor.cartesian.ConcatDeviceAndBindSchemaForPredicateVisitor;
 import org.apache.iotdb.db.queryengine.plan.expression.visitor.cartesian.ConcatExpressionWithSuffixPathsVisitor;
 import org.apache.iotdb.db.queryengine.plan.statement.component.ResultColumn;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinScalarFunction;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinTimeSeriesGeneratingFunction;
 import org.apache.iotdb.db.utils.constant.SqlConstant;
 
 import org.apache.tsfile.common.constant.TsFileConstant;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowFunctionsTask.java

@@ -24,9 +24,6 @@
 import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant;
 import org.apache.iotdb.commons.udf.UDFInformation;
 import org.apache.iotdb.commons.udf.UDFType;
-import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
-import org.apache.iotdb.commons.udf.builtin.BuiltinScalarFunction;
-import org.apache.iotdb.commons.udf.builtin.BuiltinTimeSeriesGeneratingFunction;
 import org.apache.iotdb.commons.udf.builtin.relational.TableBuiltinAggregationFunction;
 import org.apache.iotdb.commons.udf.builtin.relational.TableBuiltinScalarFunction;
 import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
@@ -35,6 +32,9 @@
 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.relational.function.TableBuiltinTableFunction;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinScalarFunction;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinTimeSeriesGeneratingFunction;
 import org.apache.iotdb.db.queryengine.plan.udf.TableUDFUtils;
 import org.apache.iotdb.db.queryengine.plan.udf.TreeUDFUtils;
 import org.apache.iotdb.rpc.TSStatusCode;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/ExpressionFactory.java

@@ -23,7 +23,6 @@
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.path.MeasurementPath;
 import org.apache.iotdb.commons.path.PartialPath;
-import org.apache.iotdb.commons.udf.builtin.BuiltinTimeSeriesGeneratingFunction;
 import org.apache.iotdb.db.queryengine.plan.expression.binary.AdditionExpression;
 import org.apache.iotdb.db.queryengine.plan.expression.binary.EqualToExpression;
 import org.apache.iotdb.db.queryengine.plan.expression.binary.GreaterEqualExpression;
@@ -46,6 +45,7 @@
 import org.apache.iotdb.db.queryengine.plan.expression.unary.LogicNotExpression;
 import org.apache.iotdb.db.queryengine.plan.expression.unary.RegularExpression;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.GroupByTimeParameter;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinTimeSeriesGeneratingFunction;
 
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.utils.TimeDuration;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/expression/multi/FunctionExpression.java

@@ -21,8 +21,6 @@
 
 import org.apache.iotdb.commons.conf.IoTDBConstant;
 import org.apache.iotdb.commons.path.PartialPath;
-import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
-import org.apache.iotdb.commons.udf.builtin.BuiltinScalarFunction;
 import org.apache.iotdb.db.queryengine.common.NodeRef;
 import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
 import org.apache.iotdb.db.queryengine.plan.expression.Expression;
@@ -31,6 +29,8 @@
 import org.apache.iotdb.db.queryengine.plan.expression.multi.builtin.BuiltInScalarFunctionHelperFactory;
 import org.apache.iotdb.db.queryengine.plan.expression.visitor.ExpressionVisitor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.InputLocation;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinScalarFunction;
 import org.apache.iotdb.db.queryengine.plan.udf.TreeUDFUtils;
 import org.apache.iotdb.db.queryengine.transformation.dag.memory.LayerMemoryAssigner;
 import org.apache.iotdb.db.queryengine.transformation.dag.udf.UDTFExecutor;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/optimization/AggregationPushDown.java

@@ -25,7 +25,6 @@
 import org.apache.iotdb.commons.path.MeasurementPath;
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant;
-import org.apache.iotdb.commons.udf.builtin.BuiltinAggregationFunction;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
 import org.apache.iotdb.db.queryengine.plan.analyze.Analysis;
@@ -57,6 +56,7 @@
 import org.apache.iotdb.db.queryengine.plan.statement.StatementType;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement;
+import org.apache.iotdb.db.queryengine.plan.udf.BuiltinAggregationFunction;
 import org.apache.iotdb.db.schemaengine.schemaregion.utils.MetaUtils;
 import org.apache.iotdb.db.utils.SchemaUtils;
 

…/builtin/BuiltinAggregationFunction.java …plan/udf/BuiltinAggregationFunction.javaiotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/builtin/BuiltinAggregationFunction.java renamed to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/BuiltinAggregationFunction.java iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/builtin/BuiltinAggregationFunction.java renamed to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/BuiltinAggregationFunction.java

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-package org.apache.iotdb.commons.udf.builtin;
+package org.apache.iotdb.db.queryengine.plan.udf;
 
 import java.util.Arrays;
 import java.util.HashSet;