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(

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java

@@ -43,8 +43,10 @@
 import org.apache.tsfile.read.common.block.column.LongColumn;
 import org.apache.tsfile.read.common.block.column.LongColumnBuilder;
 import org.apache.tsfile.read.common.block.column.RunLengthEncodedColumn;
+import org.apache.tsfile.read.common.block.column.TsBlockSerde;
 import org.apache.tsfile.utils.RamUsageEstimator;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -62,14 +64,13 @@ public class TableFunctionOperator implements ProcessOperator {
   private static final long INSTANCE_SIZE =
       RamUsageEstimator.shallowSizeOfInstance(AggregationMergeSortOperator.class);
 
-  private static final int DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES =
-      TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
-
   private final CommonOperatorContext operatorContext;
   private final Operator inputOperator;
   private final TableFunctionProcessorProvider processorProvider;
   private final PartitionRecognizer partitionRecognizer;
   private final TsBlockBuilder properBlockBuilder;
+  private final int maxTsBlockSizeInBytes;
+  private final int maxTsBlockLineNumber;
   private final int properChannelCount;
   private final boolean needPassThrough;
   private final PartitionCache partitionCache;
@@ -83,6 +84,7 @@ public class TableFunctionOperator implements ProcessOperator {
   private boolean finished = false;
 
   private final Queue<TsBlock> resultTsBlocks;
+  private final TsBlockSerde tsBlockSerde;
 
   public TableFunctionOperator(
       CommonOperatorContext operatorContext,
@@ -109,8 +111,13 @@ public TableFunctionOperator(
     this.needPassThrough = properChannelCount != outputDataTypes.size();
     this.partitionState = null;
     this.properBlockBuilder = new TsBlockBuilder(outputDataTypes.subList(0, properChannelCount));
+    this.maxTsBlockSizeInBytes =
+        TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+    this.maxTsBlockLineNumber =
+        TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber();
     this.partitionCache = new PartitionCache();
     this.resultTsBlocks = new LinkedList<>();
+    this.tsBlockSerde = new TsBlockSerde();
     this.requireRecordSnapshot = requireRecordSnapshot;
     this.ioTDBLocal = ioTDBLocal;
   }
@@ -209,7 +216,8 @@ private ColumnBuilder getPassThroughIndexBuilder() {
   }
 
   private List<TsBlock> buildTsBlock(
-      List<ColumnBuilder> properColumnBuilders, ColumnBuilder passThroughIndexBuilder) {
+      List<ColumnBuilder> properColumnBuilders, ColumnBuilder passThroughIndexBuilder)
+      throws IOException {
     int positionCount = 0;
     if (properChannelCount > 0) {
       // if there is proper column, use its position count
@@ -241,16 +249,52 @@ private List<TsBlock> buildTsBlock(
         int subBlockPositionCount = passThroughColumns[0].getPositionCount();
         TsBlock subProperBlock = properBlock.getRegion(builtCount, subBlockPositionCount);
         builtCount += subBlockPositionCount;
-        result.add(subProperBlock.appendValueColumns(passThroughColumns));
+        addSplitTsBlocks(result, subProperBlock.appendValueColumns(passThroughColumns));
       }
     } else {
-      // split the proper block into smaller blocks
-      result.add(properBlock);
+      addSplitTsBlocks(result, properBlock);
     }
     properBlockBuilder.reset();
     return result;
   }
 
+  private void addSplitTsBlocks(List<TsBlock> result, TsBlock source) throws IOException {
+    int offset = 0;
+    while (offset < source.getPositionCount()) {
+      int candidatePositionCount =
+          Math.min(source.getPositionCount() - offset, Math.max(1, maxTsBlockLineNumber));
+      int resultPositionCount =
+          getMaxSerializedPositionCount(source, offset, candidatePositionCount);
+      result.add(source.getRegion(offset, resultPositionCount));
+      offset += resultPositionCount;
+    }
+  }
+
+  private int getMaxSerializedPositionCount(TsBlock source, int offset, int candidatePositionCount)
+      throws IOException {
+    if (getSerializedSizeInBytes(source, offset, candidatePositionCount) <= maxTsBlockSizeInBytes) {
+      return candidatePositionCount;
+    }
+
+    // A row is indivisible, so return it alone if it is larger than maxTsBlockSizeInBytes.
+    int left = 1;
+    int right = candidatePositionCount - 1;
+    while (left < right) {
+      int mid = left + (right - left + 1) / 2;
+      if (getSerializedSizeInBytes(source, offset, mid) <= maxTsBlockSizeInBytes) {
+        left = mid;
+      } else {
+        right = mid - 1;
+      }
+    }
+    return left;
+  }
+
+  private int getSerializedSizeInBytes(TsBlock source, int offset, int positionCount)
+      throws IOException {
+    return tsBlockSerde.serialize(source.getRegion(offset, positionCount)).remaining();
+  }
+
   private void consumeCurrentPartitionState() {
     partitionState = null;
   }
@@ -272,6 +316,7 @@ public boolean hasNext() throws Exception {
   public void close() throws Exception {
     partitionCache.close();
     inputOperator.close();
+    resultTsBlocks.clear();
     if (processor != null) {
       destroyProcessor(processor);
       processor = null;
@@ -281,18 +326,18 @@ public void close() throws Exception {
 
   @Override
   public boolean isFinished() throws Exception {
-    return finished;
+    return finished && resultTsBlocks.isEmpty();
   }
 
   @Override
   public long calculateMaxPeekMemory() {
     return inputOperator.calculateMaxPeekMemory()
-        + Math.max(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES, properBlockBuilder.getRetainedSizeInBytes());
+        + Math.max(maxTsBlockSizeInBytes, properBlockBuilder.getRetainedSizeInBytes());
   }
 
   @Override
   public long calculateMaxReturnSize() {
-    return Math.max(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES, properBlockBuilder.getRetainedSizeInBytes());
+    return maxTsBlockSizeInBytes;
   }
 
   @Override