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 0000000000000..4a032af59016c --- /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 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 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 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 7fe6648fe70ba..9db3acb13e098 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.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 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 d4bc8a0da4c79..ee28f04c2199c 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 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 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 TableFunctionOperator( this.ioTDBLocal = ioTDBLocal; } - @Override - public CommonOperatorContext getOperatorContext() { - return this.operatorContext; - } - @Override public ListenableFuture isBlocked() { if (isBlocked == null) { @@ -148,8 +138,8 @@ private ListenableFuture tryGetNextTsBlock() { @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 TsBlock next() throws Exception { 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 TsBlock next() throws Exception { partitionCache.clear(); destroyProcessor(processor); processor = null; - return resultTsBlocks.poll(); + return getNextResultTsBlock(); } else { processor = processorProvider.getDataProcessor(); processor.beforeStart(ioTDBLocal); @@ -196,8 +186,21 @@ public TsBlock next() throws Exception { } 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 getProperColumnBuilders() { @@ -244,7 +247,6 @@ private List buildTsBlock( result.add(subProperBlock.appendValueColumns(passThroughColumns)); } } else { - // split the proper block into smaller blocks result.add(properBlock); } properBlockBuilder.reset(); @@ -265,12 +267,15 @@ private void destroyProcessor(TableFunctionDataProcessor dataProcessor) { @Override public boolean hasNext() throws Exception { - return !finished || !resultTsBlocks.isEmpty(); + return !finished || retainedTsBlock != null || !resultTsBlocks.isEmpty(); } @Override public void close() throws Exception { - partitionCache.close(); + partitionCache.clear(); + resultTsBlocks.clear(); + resultTsBlock = null; + retainedTsBlock = null; inputOperator.close(); if (processor != null) { destroyProcessor(processor); @@ -281,18 +286,18 @@ public void close() throws Exception { @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/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 bcd17aa665c47..cf1117de40ac4 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 @@ 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.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.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 void testPartitionRecognizer() { } } + @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 properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + if (!outputInFinish) { + appendRows(properColumnBuilders, passThroughIndexBuilder); + } + } + + @Override + public void finish( + List properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + if (outputInFinish) { + appendRows(properColumnBuilders, passThroughIndexBuilder); + } + } + + private void appendRows( + List 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 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 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> expected) { Iterator recordIterable = slice.getRequiredRecordIterator(false); int i = 0;