Fix table UDF result block splitting

integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java

@@ -0,0 +1,136 @@
+/*
+ * 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.query.udf.example.relational;
+
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.relational.table.MapTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.ScalarArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.ScalarParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class LargeResultTableFunction implements TableFunction {
+
+  private static final String TABLE_PARAMETER_NAME = "DATA";
+  private static final String REPEAT_COUNT_PARAMETER_NAME = "REPEAT_COUNT";
+  private static final String PAYLOAD_SIZE_PARAMETER_NAME = "PAYLOAD_SIZE";
+
+  @Override
+  public List<ParameterSpecification> getArgumentsSpecifications() {
+    return Arrays.asList(
+        TableParameterSpecification.builder()
+            .name(TABLE_PARAMETER_NAME)
+            .rowSemantics()
+            .passThroughColumns()
+            .build(),
+        ScalarParameterSpecification.builder()
+            .name(REPEAT_COUNT_PARAMETER_NAME)
+            .type(Type.INT32)
+            .build(),
+        ScalarParameterSpecification.builder()
+            .name(PAYLOAD_SIZE_PARAMETER_NAME)
+            .type(Type.INT32)
+            .build());
+  }
+
+  @Override
+  public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+    MapTableFunctionHandle handle =
+        new MapTableFunctionHandle.Builder()
+            .addProperty(
+                REPEAT_COUNT_PARAMETER_NAME,
+                ((ScalarArgument) arguments.get(REPEAT_COUNT_PARAMETER_NAME)).getValue())
+            .addProperty(
+                PAYLOAD_SIZE_PARAMETER_NAME,
+                ((ScalarArgument) arguments.get(PAYLOAD_SIZE_PARAMETER_NAME)).getValue())
+            .build();
+    return TableFunctionAnalysis.builder()
+        .properColumnSchema(
+            DescribedSchema.builder()
+                .addField("repeat_index", Type.INT32)
+                .addField("payload", Type.STRING)
+                .build())
+        .requiredColumns(TABLE_PARAMETER_NAME, Collections.singletonList(0))
+        .handle(handle)
+        .build();
+  }
+
+  @Override
+  public TableFunctionHandle createTableFunctionHandle() {
+    return new MapTableFunctionHandle();
+  }
+
+  @Override
+  public TableFunctionProcessorProvider getProcessorProvider(
+      TableFunctionHandle tableFunctionHandle) {
+    return new TableFunctionProcessorProvider() {
+      @Override
+      public TableFunctionDataProcessor getDataProcessor() {
+        return new TableFunctionDataProcessor() {
+          private final int repeatCount =
+              (int)
+                  ((MapTableFunctionHandle) tableFunctionHandle)
+                      .getProperty(REPEAT_COUNT_PARAMETER_NAME);
+          private final String payloadSuffix =
+              "x"
+                  .repeat(
+                      (int)
+                          ((MapTableFunctionHandle) tableFunctionHandle)
+                              .getProperty(PAYLOAD_SIZE_PARAMETER_NAME));
+          private long recordIndex;
+
+          @Override
+          public void process(
+              Record input,
+              List<ColumnBuilder> properColumnBuilders,
+              ColumnBuilder passThroughIndexBuilder) {
+            for (int repeatIndex = 0; repeatIndex < repeatCount; repeatIndex++) {
+              properColumnBuilders.get(0).writeInt(repeatIndex);
+              properColumnBuilders
+                  .get(1)
+                  .writeBinary(
+                      new Binary(repeatIndex + ":" + payloadSuffix, TSFileConfig.STRING_CHARSET));
+              passThroughIndexBuilder.writeLong(recordIndex);
+            }
+            recordIndex++;
+          }
+        };
+      }
+    };
+  }
+}

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java

@@ -26,13 +26,18 @@
 
 import org.junit.After;
 import org.junit.AfterClass;
+import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.HashSet;
+import java.util.Set;
 
 import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
@@ -42,6 +47,9 @@
 @Category({TableLocalStandaloneIT.class, TableClusterIT.class})
 public class IoTDBUserDefinedTableFunctionIT {
   private static final String DATABASE_NAME = "test";
+  private static final int MAX_TSBLOCK_SIZE_IN_BYTES = 1024;
+  private static final int LARGE_RESULT_REPEAT_COUNT = 64;
+  private static final int LARGE_RESULT_PAYLOAD_SIZE = 128;
   private static final String[] sqls =
       new String[] {
         "CREATE DATABASE " + DATABASE_NAME,
@@ -57,6 +65,10 @@ public class IoTDBUserDefinedTableFunctionIT {
 
   @BeforeClass
   public static void setUp() throws Exception {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getDataNodeCommonConfig()
+        .setMaxTsBlockSizeInByte(MAX_TSBLOCK_SIZE_IN_BYTES);
     EnvFactory.getEnv().initClusterEnvironment();
     insertData();
   }
@@ -182,6 +194,74 @@ public void testMyRepeat() {
         DATABASE_NAME);
   }
 
+  @Test
+  public void testLargeResultIsSplitWithoutDataLoss() throws Exception {
+    SQLFunctionUtils.createUDF(
+        "large_result",
+        "org.apache.iotdb.db.query.udf.example.relational.LargeResultTableFunction");
+
+    Set<String> returnedRows = new HashSet<>();
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE_NAME);
+      try (ResultSet resultSet =
+          statement.executeQuery(
+              "SELECT * FROM large_result(vehicle, "
+                  + LARGE_RESULT_REPEAT_COUNT
+                  + ", "
+                  + LARGE_RESULT_PAYLOAD_SIZE
+                  + ")")) {
+        while (resultSet.next()) {
+          int repeatIndex = resultSet.getInt("repeat_index");
+          long time = resultSet.getLong("time");
+          Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT);
+          Assert.assertEquals(
+              repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE),
+              resultSet.getString("payload"));
+          assertPassThroughColumns(resultSet, time);
+          Assert.assertTrue(
+              "Duplicate result row for time " + time + " and repeat index " + repeatIndex,
+              returnedRows.add(time + ":" + repeatIndex));
+        }
+      }
+    }
+
+    long[] inputTimes = new long[] {1, 2, 3, 5};
+    Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size());
+    for (long time : inputTimes) {
+      for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) {
+        Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex));
+      }
+    }
+  }
+
+  private static void assertPassThroughColumns(ResultSet resultSet, long time) throws SQLException {
+    switch ((int) time) {
+      case 1:
+        Assert.assertEquals("d0", resultSet.getString("device_id"));
+        Assert.assertEquals(1, resultSet.getInt("s1"));
+        Assert.assertEquals(1, resultSet.getLong("s2"));
+        break;
+      case 2:
+        Assert.assertEquals("d0", resultSet.getString("device_id"));
+        Assert.assertNull(resultSet.getObject("s1"));
+        Assert.assertEquals(2, resultSet.getLong("s2"));
+        break;
+      case 3:
+        Assert.assertEquals("d0", resultSet.getString("device_id"));
+        Assert.assertEquals(3, resultSet.getInt("s1"));
+        Assert.assertEquals(3, resultSet.getLong("s2"));
+        break;
+      case 5:
+        Assert.assertEquals("d1", resultSet.getString("device_id"));
+        Assert.assertEquals(4, resultSet.getInt("s1"));
+        Assert.assertNull(resultSet.getObject("s2"));
+        break;
+      default:
+        Assert.fail("Unexpected pass-through time: " + time);
+    }
+  }
+
   @Test
   public void testHybrid() {
     SQLFunctionUtils.createUDF(