This is an automated email from the ASF dual-hosted git repository.
ColinLeeo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 9fe55fce9a9 Fix table UDF result block splitting (#18333)
9fe55fce9a9 is described below
commit 9fe55fce9a933c9d4ff25d3ce10f5d8271107ca4
Author: Colin Lee <[email protected]>
AuthorDate: Thu Aug 6 11:11:05 2026 +0800
Fix table UDF result block splitting (#18333)
* Fix table UDF result block splitting
* Fix table UDF result splitting after pass-through
* Optimize table UDF result block splitting
* Estimate table UDF result blocks by memory size
* Clarify table UDF result size accounting
* Remove redundant table UDF split condition
* Reuse AbstractOperator for table UDF result splitting
* Harden table UDF result splitting lifecycle
* Remove redundant table UDF type list copy
* Clear partition cache on close
---
.../relational/LargeResultTableFunction.java | 136 +++++++++
.../db/it/udf/IoTDBUserDefinedTableFunctionIT.java | 80 +++++
.../process/function/TableFunctionOperator.java | 51 ++--
.../process/function/partition/PartitionCache.java | 2 +-
.../process/tvf/TableFunctionOperatorTest.java | 331 +++++++++++++++++++++
5 files changed, 576 insertions(+), 24 deletions(-)
diff --git
a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java
b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java
new file mode 100644
index 00000000000..4a032af5901
--- /dev/null
+++
b/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++;
+ }
+ };
+ }
+ };
+ }
+}
diff --git
a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java
b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java
index 7fe6648fe70..9db3acb13e0 100644
---
a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java
+++
b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java
@@ -26,13 +26,18 @@ import
org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
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 @@ import static org.junit.Assert.fail;
@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 class IoTDBUserDefinedTableFunctionIT {
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(
diff --git
a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
index d4bc8a0da4c..fb9b90c9f60 100644
---
a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
+++
b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
@@ -19,9 +19,9 @@
package org.apache.iotdb.calc.execution.operator.process.function;
+import org.apache.iotdb.calc.execution.operator.AbstractOperator;
import org.apache.iotdb.calc.execution.operator.CommonOperatorContext;
import org.apache.iotdb.calc.execution.operator.Operator;
-import
org.apache.iotdb.calc.execution.operator.process.AggregationMergeSortOperator;
import org.apache.iotdb.calc.execution.operator.process.ProcessOperator;
import
org.apache.iotdb.calc.execution.operator.process.function.partition.PartitionCache;
import
org.apache.iotdb.calc.execution.operator.process.function.partition.PartitionState;
@@ -36,7 +36,6 @@ import
org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProc
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.common.conf.TSFileDescriptor;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.read.common.block.TsBlock;
import org.apache.tsfile.read.common.block.TsBlockBuilder;
@@ -57,15 +56,11 @@ import java.util.Queue;
import static com.google.common.base.Preconditions.checkArgument;
// only one input source is supported now
-public class TableFunctionOperator implements ProcessOperator {
+public class TableFunctionOperator extends AbstractOperator implements
ProcessOperator {
private static final long INSTANCE_SIZE =
-
RamUsageEstimator.shallowSizeOfInstance(AggregationMergeSortOperator.class);
+ RamUsageEstimator.shallowSizeOfInstance(TableFunctionOperator.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;
@@ -115,11 +110,6 @@ public class TableFunctionOperator implements
ProcessOperator {
this.ioTDBLocal = ioTDBLocal;
}
- @Override
- public CommonOperatorContext getOperatorContext() {
- return this.operatorContext;
- }
-
@Override
public ListenableFuture<?> isBlocked() {
if (isBlocked == null) {
@@ -148,8 +138,8 @@ public class TableFunctionOperator implements
ProcessOperator {
@Override
public TsBlock next() throws Exception {
- if (!resultTsBlocks.isEmpty()) {
- return resultTsBlocks.poll();
+ if (retainedTsBlock != null || !resultTsBlocks.isEmpty()) {
+ return getNextResultTsBlock();
}
if (partitionState == null) {
partitionState = partitionRecognizer.nextState();
@@ -172,7 +162,7 @@ public class TableFunctionOperator implements
ProcessOperator {
resultTsBlocks.addAll(buildTsBlock(properColumnBuilders,
passThroughIndexBuilder));
partitionCache.clear();
consumeCurrentPartitionState();
- return resultTsBlocks.poll();
+ return getNextResultTsBlock();
}
if (stateType == PartitionState.StateType.NEW_PARTITION) {
if (processor != null) {
@@ -182,7 +172,7 @@ public class TableFunctionOperator implements
ProcessOperator {
partitionCache.clear();
destroyProcessor(processor);
processor = null;
- return resultTsBlocks.poll();
+ return getNextResultTsBlock();
} else {
processor = processorProvider.getDataProcessor();
processor.beforeStart(ioTDBLocal);
@@ -196,8 +186,21 @@ public class TableFunctionOperator implements
ProcessOperator {
}
consumeCurrentPartitionState();
resultTsBlocks.addAll(buildTsBlock(properColumnBuilders,
passThroughIndexBuilder));
- return resultTsBlocks.poll();
+ return getNextResultTsBlock();
+ }
+ }
+
+ /**
+ * Applies {@link AbstractOperator}'s low-cost row-count splitting after
pass-through columns have
+ * been appended. The configured byte size is a logical target rather than
an exact serialized
+ * limit; in particular, a single oversized row and highly variable binary
payloads may exceed it.
+ */
+ private TsBlock getNextResultTsBlock() {
+ if (retainedTsBlock != null) {
+ return getResultFromRetainedTsBlock();
}
+ resultTsBlock = resultTsBlocks.poll();
+ return resultTsBlock == null ? null : checkTsBlockSizeAndGetResult();
}
private List<ColumnBuilder> getProperColumnBuilders() {
@@ -244,7 +247,6 @@ public class TableFunctionOperator implements
ProcessOperator {
result.add(subProperBlock.appendValueColumns(passThroughColumns));
}
} else {
- // split the proper block into smaller blocks
result.add(properBlock);
}
properBlockBuilder.reset();
@@ -265,12 +267,15 @@ public class TableFunctionOperator implements
ProcessOperator {
@Override
public boolean hasNext() throws Exception {
- return !finished || !resultTsBlocks.isEmpty();
+ return !finished || retainedTsBlock != null || !resultTsBlocks.isEmpty();
}
@Override
public void close() throws Exception {
partitionCache.close();
+ resultTsBlocks.clear();
+ resultTsBlock = null;
+ retainedTsBlock = null;
inputOperator.close();
if (processor != null) {
destroyProcessor(processor);
@@ -281,18 +286,18 @@ public class TableFunctionOperator implements
ProcessOperator {
@Override
public boolean isFinished() throws Exception {
- return finished;
+ return finished && retainedTsBlock == null && resultTsBlocks.isEmpty();
}
@Override
public long calculateMaxPeekMemory() {
return inputOperator.calculateMaxPeekMemory()
- + Math.max(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES,
properBlockBuilder.getRetainedSizeInBytes());
+ + Math.max(maxReturnSize, properBlockBuilder.getRetainedSizeInBytes());
}
@Override
public long calculateMaxReturnSize() {
- return Math.max(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES,
properBlockBuilder.getRetainedSizeInBytes());
+ return maxReturnSize;
}
@Override
diff --git
a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/partition/PartitionCache.java
b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/partition/PartitionCache.java
index a33113d8f48..c0f6d64415c 100644
---
a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/partition/PartitionCache.java
+++
b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/partition/PartitionCache.java
@@ -117,6 +117,6 @@ public class PartitionCache {
}
public void close() {
- // do nothing
+ clear();
}
}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/process/tvf/TableFunctionOperatorTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/process/tvf/TableFunctionOperatorTest.java
index bcd17aa665c..cf1117de40a 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/process/tvf/TableFunctionOperatorTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/process/tvf/TableFunctionOperatorTest.java
@@ -21,9 +21,11 @@ package
org.apache.iotdb.db.queryengine.execution.operator.process.tvf;
import org.apache.iotdb.calc.execution.operator.Operator;
import
org.apache.iotdb.calc.execution.operator.process.function.PartitionRecognizer;
+import
org.apache.iotdb.calc.execution.operator.process.function.TableFunctionOperator;
import
org.apache.iotdb.calc.execution.operator.process.function.partition.PartitionState;
import
org.apache.iotdb.calc.execution.operator.process.function.partition.Slice;
import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
import org.apache.iotdb.db.queryengine.common.QueryId;
@@ -31,9 +33,14 @@ import
org.apache.iotdb.db.queryengine.execution.driver.DriverContext;
import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.relational.access.Record;
+import
org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import
org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor;
+import org.apache.tsfile.block.column.ColumnBuilder;
import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.common.conf.TSFileDescriptor;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.read.common.block.TsBlock;
import org.apache.tsfile.read.common.block.TsBlockBuilder;
@@ -48,11 +55,15 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
import static
org.apache.iotdb.calc.plan.planner.CommonOperatorUtils.TIME_COLUMN_TEMPLATE;
import static
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
public class TableFunctionOperatorTest {
private static final ExecutorService instanceNotificationExecutor =
@@ -287,6 +298,326 @@ public class TableFunctionOperatorTest {
}
}
+ @Test
+ public void testResultTsBlockUsesAbstractOperatorSplitting() throws
Exception {
+ assertResultTsBlockUsesAbstractOperatorSplitting(false, false);
+ assertResultTsBlockUsesAbstractOperatorSplitting(true, false);
+ }
+
+ @Test
+ public void testFinishResultUsesAbstractOperatorSplitting() throws Exception
{
+ assertResultTsBlockUsesAbstractOperatorSplitting(false, true);
+ assertResultTsBlockUsesAbstractOperatorSplitting(true, true);
+ }
+
+ private void assertResultTsBlockUsesAbstractOperatorSplitting(
+ boolean withPassThrough, boolean outputInFinish) throws Exception {
+ int originalMaxBlockSize =
+ TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+ int maxBlockSize = 128;
+ // AbstractOperator uses at least one byte as the estimated size of each
row, so producing more
+ // rows than the configured byte limit guarantees that the result must be
split.
+ int outputRowCount = maxBlockSize * 2;
+ try {
+ OperatorContext operatorContext =
+ createOperatorContext(
+ "abstract_operator_split_" + withPassThrough + "_" +
outputInFinish);
+ TableFunctionProcessorProvider provider =
+ new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionDataProcessor getDataProcessor() {
+ return new TableFunctionDataProcessor() {
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ if (!outputInFinish) {
+ appendRows(properColumnBuilders, passThroughIndexBuilder);
+ }
+ }
+
+ @Override
+ public void finish(
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ if (outputInFinish) {
+ appendRows(properColumnBuilders, passThroughIndexBuilder);
+ }
+ }
+
+ private void appendRows(
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ for (int i = 0; i < outputRowCount; i++) {
+ properColumnBuilders.get(0).writeLong(i);
+ if (passThroughIndexBuilder != null) {
+ passThroughIndexBuilder.writeLong(0);
+ }
+ }
+ }
+ };
+ }
+ };
+
+ Operator singleRowChild = constructLongChildOperator(operatorContext,
new long[] {1});
+
+ // FragmentInstanceContext initialization reloads the TsFile
configuration, so apply the test
+ // limit immediately before constructing the operator that captures it.
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(maxBlockSize);
+ int returnedRows = 0;
+ int returnedBlocks = 0;
+ try (TableFunctionOperator operator =
+ new TableFunctionOperator(
+ operatorContext,
+ provider,
+ singleRowChild,
+ Collections.singletonList(TSDataType.INT64),
+ withPassThrough
+ ? Arrays.asList(TSDataType.INT64, TSDataType.INT64)
+ : Collections.singletonList(TSDataType.INT64),
+ 1,
+ Collections.singletonList(0),
+ withPassThrough ? Collections.singletonList(0) :
Collections.emptyList(),
+ withPassThrough,
+ Collections.emptyList(),
+ false,
+ mock(IoTDBLocal.class))) {
+ assertEquals(maxBlockSize, operator.calculateMaxReturnSize());
+ while (!operator.isFinished()) {
+ operator.isBlocked();
+ TsBlock block = operator.next();
+ if (block == null) {
+ continue;
+ }
+ returnedBlocks++;
+ for (int i = 0; i < block.getPositionCount(); i++) {
+ assertEquals(returnedRows, block.getColumn(0).getLong(i));
+ if (withPassThrough) {
+ assertEquals(1, block.getColumn(1).getLong(i));
+ }
+ returnedRows++;
+ }
+ }
+ }
+
+ assertEquals(outputRowCount, returnedRows);
+ Assert.assertTrue("Returned block count: " + returnedBlocks,
returnedBlocks > 1);
+ } finally {
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(originalMaxBlockSize);
+ }
+ }
+
+ @Test
+ public void testPartitionResultsAreDrainedBeforeProcessingNextPartition()
throws Exception {
+ int originalMaxBlockSize =
+ TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+ int maxBlockSize = 128;
+ int outputRowCount = maxBlockSize * 2;
+ try {
+ OperatorContext operatorContext =
createOperatorContext("abstract_operator_split_partitions");
+ AtomicInteger processorCount = new AtomicInteger();
+ TableFunctionProcessorProvider provider =
+ new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionDataProcessor getDataProcessor() {
+ int resultOffset = processorCount.getAndIncrement() *
outputRowCount;
+ return new TableFunctionDataProcessor() {
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ for (int i = 0; i < outputRowCount; i++) {
+ properColumnBuilders.get(0).writeLong(resultOffset + i);
+ }
+ }
+ };
+ }
+ };
+ Operator twoPartitionChild = constructLongChildOperator(operatorContext,
new long[] {1, 2});
+
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(maxBlockSize);
+ int returnedRows = 0;
+ int returnedBlocks = 0;
+ try (TableFunctionOperator operator =
+ new TableFunctionOperator(
+ operatorContext,
+ provider,
+ twoPartitionChild,
+ Collections.singletonList(TSDataType.INT64),
+ Collections.singletonList(TSDataType.INT64),
+ 1,
+ Collections.singletonList(0),
+ Collections.emptyList(),
+ false,
+ Collections.singletonList(0),
+ false,
+ mock(IoTDBLocal.class))) {
+ while (!operator.isFinished()) {
+ operator.isBlocked();
+ TsBlock block = operator.next();
+ if (block == null) {
+ continue;
+ }
+ returnedBlocks++;
+ for (int i = 0; i < block.getPositionCount(); i++) {
+ assertEquals(returnedRows++, block.getColumn(0).getLong(i));
+ }
+ }
+ }
+
+ assertEquals(2, processorCount.get());
+ assertEquals(2 * outputRowCount, returnedRows);
+ Assert.assertTrue("Returned block count: " + returnedBlocks,
returnedBlocks > 2);
+ } finally {
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(originalMaxBlockSize);
+ }
+ }
+
+ @Test
+ public void testCloseReleasesPendingResultState() throws Exception {
+ int originalMaxBlockSize =
+ TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+ int maxBlockSize = 128;
+ int outputRowCount = maxBlockSize * 2;
+ try {
+ OperatorContext operatorContext =
+ createOperatorContext("close_pending_table_function_result");
+ AtomicBoolean processorDestroyed = new AtomicBoolean();
+ TableFunctionProcessorProvider provider =
+ new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionDataProcessor getDataProcessor() {
+ return new TableFunctionDataProcessor() {
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ for (int i = 0; i < outputRowCount; i++) {
+ properColumnBuilders.get(0).writeLong(i);
+ }
+ }
+
+ @Override
+ public void beforeDestroy() {
+ processorDestroyed.set(true);
+ }
+ };
+ }
+ };
+ AtomicBoolean childClosed = new AtomicBoolean();
+ Operator child = constructLongChildOperator(operatorContext, new long[]
{1}, childClosed);
+ IoTDBLocal ioTDBLocal = mock(IoTDBLocal.class);
+
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(maxBlockSize);
+ TableFunctionOperator operator =
+ new TableFunctionOperator(
+ operatorContext,
+ provider,
+ child,
+ Collections.singletonList(TSDataType.INT64),
+ Collections.singletonList(TSDataType.INT64),
+ 1,
+ Collections.singletonList(0),
+ Collections.emptyList(),
+ false,
+ Collections.emptyList(),
+ false,
+ ioTDBLocal);
+ operator.isBlocked();
+ Assert.assertNotNull(operator.next());
+ long retainedSizeBeforeClose = operator.ramBytesUsed();
+
+ operator.close();
+
+ Assert.assertTrue(childClosed.get());
+ Assert.assertTrue(processorDestroyed.get());
+ verify(ioTDBLocal).close();
+ Assert.assertTrue(operator.ramBytesUsed() < retainedSizeBeforeClose);
+ } finally {
+
TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(originalMaxBlockSize);
+ }
+ }
+
+ private Operator constructLongChildOperator(
+ OperatorContext operatorContext, long[] values, AtomicBoolean closed) {
+ return new Operator() {
+ private boolean consumed;
+
+ @Override
+ public OperatorContext getOperatorContext() {
+ return operatorContext;
+ }
+
+ @Override
+ public TsBlock next() {
+ TsBlockBuilder builder =
+ new TsBlockBuilder(values.length,
Collections.singletonList(TSDataType.INT64));
+ for (long value : values) {
+ builder.getColumnBuilder(0).writeLong(value);
+ builder.declarePosition();
+ }
+ consumed = true;
+ return builder.build(
+ new RunLengthEncodedColumn(TIME_COLUMN_TEMPLATE,
builder.getPositionCount()));
+ }
+
+ @Override
+ public boolean hasNext() {
+ return !consumed;
+ }
+
+ @Override
+ public void close() {
+ closed.set(true);
+ }
+
+ @Override
+ public boolean isFinished() {
+ return consumed;
+ }
+
+ @Override
+ public long calculateMaxPeekMemory() {
+ return 0;
+ }
+
+ @Override
+ public long calculateMaxReturnSize() {
+ return 0;
+ }
+
+ @Override
+ public long calculateRetainedSizeAfterCallingNext() {
+ return 0;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0;
+ }
+ };
+ }
+
+ private Operator constructLongChildOperator(OperatorContext operatorContext,
long[] values) {
+ return constructLongChildOperator(operatorContext, values, new
AtomicBoolean());
+ }
+
+ private OperatorContext createOperatorContext(String queryName) {
+ QueryId queryId = new QueryId(queryName);
+ FragmentInstanceId instanceId =
+ new FragmentInstanceId(new PlanFragmentId(queryId, 0),
"stub-instance");
+ FragmentInstanceStateMachine stateMachine =
+ new FragmentInstanceStateMachine(instanceId,
instanceNotificationExecutor);
+ FragmentInstanceContext fragmentInstanceContext =
+ createFragmentInstanceContext(instanceId, stateMachine);
+ DriverContext driverContext = new DriverContext(fragmentInstanceContext,
0);
+ return driverContext.addOperatorContext(
+ 0, new PlanNodeId("tvf"), TableFunctionOperator.class.getSimpleName());
+ }
+
private void checkIteratorSimply(Slice slice, List<List<Object>> expected) {
Iterator<Record> recordIterable = slice.getRequiredRecordIterator(false);
int i = 0;
