You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by co...@apache.org on 2026/07/24 13:11:54 UTC
(iotdb) branch codex/fix-table-udf-result-chunking created (now 82c40ee3a9a)
This is an automated email from the ASF dual-hosted git repository.
ColinLeeo pushed a change to branch codex/fix-table-udf-result-chunking
in repository https://gitbox.apache.org/repos/asf/iotdb.git
at 82c40ee3a9a Fix table UDF result block splitting
This branch includes the following new commits:
new 82c40ee3a9a Fix table UDF result block splitting
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(iotdb) 01/01: Fix table UDF result block splitting
Posted by co...@apache.org.
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 <sh...@163.com>
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;