This is an automated email from the ASF dual-hosted git repository. ColinLeeo pushed a commit to branch codex/fix-table-udf-result-chunking in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 82c40ee3a9af8b65e9aca7317cc2319c5e80cceb Author: ColinLee <[email protected]> AuthorDate: Fri Jul 24 21:11:33 2026 +0800 Fix table UDF result block splitting --- .../process/function/TableFunctionOperator.java | 89 ++++++++++-- .../process/tvf/TableFunctionOperatorTest.java | 150 +++++++++++++++++++++ 2 files changed, 229 insertions(+), 10 deletions(-) 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..56fd9b4843b 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,6 +19,7 @@ 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; @@ -57,7 +58,7 @@ 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); @@ -65,11 +66,11 @@ public class TableFunctionOperator implements ProcessOperator { 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 maxTsBlockLineNumber; private final int properChannelCount; private final boolean needPassThrough; private final PartitionCache partitionCache; @@ -109,6 +110,8 @@ public class TableFunctionOperator implements ProcessOperator { this.needPassThrough = properChannelCount != outputDataTypes.size(); this.partitionState = null; this.properBlockBuilder = new TsBlockBuilder(outputDataTypes.subList(0, properChannelCount)); + this.maxTsBlockLineNumber = + TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); this.requireRecordSnapshot = requireRecordSnapshot; @@ -148,8 +151,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 +175,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 +185,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,10 +199,72 @@ public class TableFunctionOperator implements ProcessOperator { } consumeCurrentPartitionState(); resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder)); - return resultTsBlocks.poll(); + return getNextResultTsBlock(); } } + private TsBlock getNextResultTsBlock() { + if (retainedTsBlock == null) { + retainedTsBlock = resultTsBlocks.poll(); + startOffset = 0; + } + return retainedTsBlock == null ? null : getResultFromRetainedTsBlock(); + } + + @Override + public TsBlock getResultFromRetainedTsBlock() { + int remainingPositionCount = retainedTsBlock.getPositionCount() - startOffset; + int candidatePositionCount = + Math.min(remainingPositionCount, Math.max(1, maxTsBlockLineNumber)); + int resultPositionCount = getMaxResultPositionCount(candidatePositionCount); + TsBlock result = + startOffset == 0 && resultPositionCount == retainedTsBlock.getPositionCount() + ? retainedTsBlock + : copyTsBlockRegion(retainedTsBlock, startOffset, resultPositionCount); + startOffset += resultPositionCount; + if (startOffset == retainedTsBlock.getPositionCount()) { + retainedTsBlock = null; + startOffset = 0; + } + return result; + } + + private int getMaxResultPositionCount(int candidatePositionCount) { + if (getRegionSizeInBytes(candidatePositionCount) <= maxReturnSize) { + return candidatePositionCount; + } + + // A row is indivisible, so keep the existing one-row-at-a-time fallback when a single row is + // larger than maxReturnSize. + int left = 1; + int right = candidatePositionCount - 1; + while (left < right) { + int mid = left + (right - left + 1) / 2; + if (getRegionSizeInBytes(mid) <= maxReturnSize) { + left = mid; + } else { + right = mid - 1; + } + } + return left; + } + + private long getRegionSizeInBytes(int positionCount) { + // getRegion() keeps the source column's backing arrays and estimates a region's size from their + // capacity. Copying the region first makes variable-width columns report the logical size of + // exactly the rows being considered. + return copyTsBlockRegion(retainedTsBlock, startOffset, positionCount).getSizeInBytes(); + } + + private TsBlock copyTsBlockRegion(TsBlock source, int offset, int positionCount) { + Column[] valueColumns = new Column[source.getValueColumnCount()]; + for (int i = 0; i < valueColumns.length; i++) { + valueColumns[i] = source.getColumn(i).getRegionCopy(offset, positionCount); + } + return new TsBlock( + positionCount, source.getTimeColumn().getRegionCopy(offset, positionCount), valueColumns); + } + private List<ColumnBuilder> getProperColumnBuilders() { return Arrays.asList(properBlockBuilder.getValueColumnBuilders()); } @@ -265,13 +330,17 @@ 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(); inputOperator.close(); + resultTsBlocks.clear(); + resultTsBlock = null; + retainedTsBlock = null; + startOffset = 0; if (processor != null) { destroyProcessor(processor); processor = null; @@ -281,7 +350,7 @@ public class TableFunctionOperator implements ProcessOperator { @Override public boolean isFinished() throws Exception { - return finished; + return finished && retainedTsBlock == null && resultTsBlocks.isEmpty(); } @Override @@ -292,7 +361,7 @@ public class TableFunctionOperator implements ProcessOperator { @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 bcd17aa665c..de57720e9a5 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,13 +33,19 @@ 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; import org.apache.tsfile.read.common.block.column.RunLengthEncodedColumn; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.apache.tsfile.utils.Binary; import org.junit.AfterClass; import org.junit.Assert; @@ -53,6 +61,7 @@ import static org.apache.iotdb.calc.plan.planner.CommonOperatorUtils.TIME_COLUMN 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; public class TableFunctionOperatorTest { private static final ExecutorService instanceNotificationExecutor = @@ -287,6 +296,147 @@ public class TableFunctionOperatorTest { } } + @Test + public void testVariableWidthResultsAreSplitByActualSize() throws Exception { + QueryId queryId = new QueryId("large_finish_result"); + 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); + OperatorContext operatorContext = + driverContext.addOperatorContext( + 0, new PlanNodeId("tvf"), TableFunctionOperator.class.getSimpleName()); + int maxLineNumber = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); + int maxBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + Assert.assertTrue(maxLineNumber >= 3); + int wideRowCount = maxLineNumber + 1; + int wideValueLength = maxBlockSize / (maxLineNumber - 1) + 1; + Binary narrowValue = new Binary("narrow", TSFileConfig.STRING_CHARSET); + Binary wideValue = new Binary(new byte[wideValueLength]); + TableFunctionProcessorProvider provider = + new TableFunctionProcessorProvider() { + @Override + public TableFunctionDataProcessor getDataProcessor() { + return new TableFunctionDataProcessor() { + @Override + public void process( + Record input, + List<ColumnBuilder> properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + // Initialize the splitter with a narrow result block. + properColumnBuilders.get(0).writeBinary(narrowValue); + } + + @Override + public void finish( + List<ColumnBuilder> properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { + // This second result block exceeds both maxLineNumber and maxBlockSize. Reusing a + // row-count estimate from the narrow block must not let an oversized slice pass + // through. + for (int i = 0; i < wideRowCount; i++) { + properColumnBuilders.get(0).writeBinary(wideValue); + } + } + }; + } + }; + + Operator singleRowChild = + new Operator() { + private boolean consumed; + + @Override + public OperatorContext getOperatorContext() { + return operatorContext; + } + + @Override + public TsBlock next() { + TsBlockBuilder builder = + new TsBlockBuilder(1, Collections.singletonList(TSDataType.INT64)); + builder.getColumnBuilder(0).writeLong(1); + builder.declarePosition(); + consumed = true; + return builder.build( + new RunLengthEncodedColumn(TIME_COLUMN_TEMPLATE, builder.getPositionCount())); + } + + @Override + public boolean hasNext() { + return !consumed; + } + + @Override + public void close() {} + + @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; + } + }; + + int returnedRows = 0; + int returnedBlocks = 0; + try (TableFunctionOperator operator = + new TableFunctionOperator( + operatorContext, + provider, + singleRowChild, + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(TSDataType.TEXT), + 1, + Collections.singletonList(0), + Collections.emptyList(), + false, + Collections.emptyList(), + false, + mock(IoTDBLocal.class))) { + while (!operator.isFinished()) { + operator.isBlocked(); + TsBlock block = operator.next(); + if (block == null) { + continue; + } + returnedBlocks++; + returnedRows += block.getPositionCount(); + Assert.assertTrue(block.getPositionCount() <= maxLineNumber); + Assert.assertTrue(block.getSizeInBytes() <= maxBlockSize); + Assert.assertTrue(new TsBlockSerde().serialize(block).remaining() <= maxBlockSize); + for (int i = 0; i < block.getPositionCount(); i++) { + int expectedLength = + returnedRows - block.getPositionCount() + i == 0 ? 6 : wideValueLength; + assertEquals(expectedLength, block.getColumn(0).getBinary(i).getLength()); + } + } + } + + assertEquals(wideRowCount + 1, returnedRows); + Assert.assertTrue(returnedBlocks > 2); + } + private void checkIteratorSimply(Slice slice, List<List<Object>> expected) { Iterator<Record> recordIterable = slice.getRequiredRecordIterator(false); int i = 0;
