From 82c40ee3a9af8b65e9aca7317cc2319c5e80cceb Mon Sep 17 00:00:00 2001 From: ColinLee Date: Fri, 24 Jul 2026 21:11:33 +0800 Subject: [PATCH 1/6] Fix table UDF result block splitting --- .../function/TableFunctionOperator.java | 89 +++++++++-- .../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 d4bc8a0da4c79..56fd9b4843b5e 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 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 TableFunctionOperator( 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 @@ 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 +175,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 +185,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,10 +199,72 @@ public TsBlock next() throws Exception { } 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 getProperColumnBuilders() { return Arrays.asList(properBlockBuilder.getValueColumnBuilders()); } @@ -265,13 +330,17 @@ 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(); inputOperator.close(); + resultTsBlocks.clear(); + resultTsBlock = null; + retainedTsBlock = null; + startOffset = 0; if (processor != null) { destroyProcessor(processor); processor = null; @@ -281,7 +350,7 @@ public void close() throws Exception { @Override public boolean isFinished() throws Exception { - return finished; + return finished && retainedTsBlock == null && resultTsBlocks.isEmpty(); } @Override @@ -292,7 +361,7 @@ public long calculateMaxPeekMemory() { @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..de57720e9a505 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,13 +33,19 @@ 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.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 void testPartitionRecognizer() { } } + @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 properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + // Initialize the splitter with a narrow result block. + properColumnBuilders.get(0).writeBinary(narrowValue); + } + + @Override + public void finish( + List 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> expected) { Iterator recordIterable = slice.getRequiredRecordIterator(false); int i = 0; From 0d46ba8d5926c35eb4fb0a5fa3162c9369c0237b Mon Sep 17 00:00:00 2001 From: ColinLee Date: Tue, 28 Jul 2026 16:12:24 +0800 Subject: [PATCH 2/6] Fix table UDF result splitting after pass-through --- .../relational/LargeResultTableFunction.java | 136 +++++++++++++++++ .../udf/IoTDBUserDefinedTableFunctionIT.java | 80 ++++++++++ .../function/TableFunctionOperator.java | 142 ++++++++---------- .../tvf/TableFunctionOperatorTest.java | 35 +++-- 4 files changed, 300 insertions(+), 93 deletions(-) create mode 100644 integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java 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 56fd9b4843b5e..1a2ae84601fb6 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,7 +19,6 @@ 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; @@ -44,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; @@ -58,18 +59,17 @@ import static com.google.common.base.Preconditions.checkArgument; // only one input source is supported now -public class TableFunctionOperator extends AbstractOperator implements ProcessOperator { +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; @@ -84,6 +84,7 @@ public class TableFunctionOperator extends AbstractOperator implements ProcessOp private boolean finished = false; private final Queue resultTsBlocks; + private final TsBlockSerde tsBlockSerde; public TableFunctionOperator( CommonOperatorContext operatorContext, @@ -110,10 +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; } @@ -151,8 +155,8 @@ private ListenableFuture tryGetNextTsBlock() { @Override public TsBlock next() throws Exception { - if (retainedTsBlock != null || !resultTsBlocks.isEmpty()) { - return getNextResultTsBlock(); + if (!resultTsBlocks.isEmpty()) { + return resultTsBlocks.poll(); } if (partitionState == null) { partitionState = partitionRecognizer.nextState(); @@ -175,7 +179,7 @@ public TsBlock next() throws Exception { resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder)); partitionCache.clear(); consumeCurrentPartitionState(); - return getNextResultTsBlock(); + return resultTsBlocks.poll(); } if (stateType == PartitionState.StateType.NEW_PARTITION) { if (processor != null) { @@ -185,7 +189,7 @@ public TsBlock next() throws Exception { partitionCache.clear(); destroyProcessor(processor); processor = null; - return getNextResultTsBlock(); + return resultTsBlocks.poll(); } else { processor = processorProvider.getDataProcessor(); processor.beforeStart(ioTDBLocal); @@ -199,70 +203,8 @@ public TsBlock next() throws Exception { } consumeCurrentPartitionState(); resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder)); - return getNextResultTsBlock(); - } - } - - private TsBlock getNextResultTsBlock() { - if (retainedTsBlock == null) { - retainedTsBlock = resultTsBlocks.poll(); - startOffset = 0; + return resultTsBlocks.poll(); } - 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 getProperColumnBuilders() { @@ -274,7 +216,8 @@ private ColumnBuilder getPassThroughIndexBuilder() { } private List buildTsBlock( - List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { + List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) + throws IOException { int positionCount = 0; if (properChannelCount > 0) { // if there is proper column, use its position count @@ -306,16 +249,52 @@ private List 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 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; } @@ -330,7 +309,7 @@ private void destroyProcessor(TableFunctionDataProcessor dataProcessor) { @Override public boolean hasNext() throws Exception { - return !finished || retainedTsBlock != null || !resultTsBlocks.isEmpty(); + return !finished || !resultTsBlocks.isEmpty(); } @Override @@ -338,9 +317,6 @@ public void close() throws Exception { partitionCache.close(); inputOperator.close(); resultTsBlocks.clear(); - resultTsBlock = null; - retainedTsBlock = null; - startOffset = 0; if (processor != null) { destroyProcessor(processor); processor = null; @@ -350,18 +326,18 @@ public void close() throws Exception { @Override public boolean isFinished() throws Exception { - return finished && retainedTsBlock == null && resultTsBlocks.isEmpty(); + 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 maxReturnSize; + return maxTsBlockSizeInBytes; } @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 de57720e9a505..bc10de7e6abd5 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 @@ -297,8 +297,14 @@ public void testPartitionRecognizer() { } @Test - public void testVariableWidthResultsAreSplitByActualSize() throws Exception { - QueryId queryId = new QueryId("large_finish_result"); + public void testVariableWidthResultsAreSplitBySerializedSize() throws Exception { + assertVariableWidthResultsAreSplitBySerializedSize(false); + assertVariableWidthResultsAreSplitBySerializedSize(true); + } + + private void assertVariableWidthResultsAreSplitBySerializedSize(boolean withPassThrough) + throws Exception { + QueryId queryId = new QueryId("large_finish_result_" + withPassThrough); FragmentInstanceId instanceId = new FragmentInstanceId(new PlanFragmentId(queryId, 0), "stub-instance"); FragmentInstanceStateMachine stateMachine = @@ -326,18 +332,23 @@ public void process( Record input, List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { - // Initialize the splitter with a narrow result block. + // Produce a narrow block before the wide result returned by finish(). properColumnBuilders.get(0).writeBinary(narrowValue); + if (passThroughIndexBuilder != null) { + passThroughIndexBuilder.writeLong(0); + } } @Override public void finish( List 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. + // This result exceeds both limits and must be split after pass-through columns are + // appended. for (int i = 0; i < wideRowCount; i++) { properColumnBuilders.get(0).writeBinary(wideValue); + if (passThroughIndexBuilder != null) { + passThroughIndexBuilder.writeLong(0); + } } } }; @@ -406,11 +417,13 @@ public long ramBytesUsed() { provider, singleRowChild, Collections.singletonList(TSDataType.INT64), - Collections.singletonList(TSDataType.TEXT), + withPassThrough + ? Arrays.asList(TSDataType.TEXT, TSDataType.INT64) + : Collections.singletonList(TSDataType.TEXT), 1, Collections.singletonList(0), - Collections.emptyList(), - false, + withPassThrough ? Collections.singletonList(0) : Collections.emptyList(), + withPassThrough, Collections.emptyList(), false, mock(IoTDBLocal.class))) { @@ -423,12 +436,14 @@ public long ramBytesUsed() { 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()); + if (withPassThrough) { + assertEquals(1, block.getColumn(1).getLong(i)); + } } } } From 6f7c588e50f89a537a4f8bab614f1b8fcb623bb5 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 30 Jul 2026 12:00:38 +0800 Subject: [PATCH 3/6] Optimize table UDF result block splitting --- .../function/TableFunctionOperator.java | 161 ++++++++++++++---- .../tvf/TableFunctionOperatorTest.java | 23 ++- 2 files changed, 141 insertions(+), 43 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 1a2ae84601fb6..289cf63b49423 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 @@ -43,10 +43,8 @@ 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; @@ -63,6 +61,13 @@ public class TableFunctionOperator implements ProcessOperator { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(AggregationMergeSortOperator.class); + // Allow for the null marker, nested encoding metadata, and dictionary metadata of each column. + // These values deliberately overestimate array-encoded columns so the result can be sliced + // without serializing it first. + private static final int ESTIMATED_SERIALIZED_COLUMN_METADATA_SIZE_IN_BYTES = + 2 * Integer.BYTES + 2 * Byte.BYTES; + private static final int ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES = + Byte.BYTES + Integer.BYTES; private final CommonOperatorContext operatorContext; private final Operator inputOperator; @@ -71,6 +76,9 @@ public class TableFunctionOperator implements ProcessOperator { private final TsBlockBuilder properBlockBuilder; private final int maxTsBlockSizeInBytes; private final int maxTsBlockLineNumber; + private final long estimatedSerializedBlockFixedSizeInBytes; + private final long estimatedSerializedPositionFixedSizeInBytes; + private final int[] variableWidthColumnIndexes; private final int properChannelCount; private final boolean needPassThrough; private final PartitionCache partitionCache; @@ -84,7 +92,6 @@ public class TableFunctionOperator implements ProcessOperator { private boolean finished = false; private final Queue resultTsBlocks; - private final TsBlockSerde tsBlockSerde; public TableFunctionOperator( CommonOperatorContext operatorContext, @@ -107,17 +114,22 @@ public TableFunctionOperator( this.partitionRecognizer = new PartitionRecognizer( partitionChannels, requiredChannels, passThroughChannels, inputDataTypes); + List resultDataTypes = new ArrayList<>(outputDataTypes); this.isDeclaredAsPassThrough = isDeclaredAsPassThrough; - this.needPassThrough = properChannelCount != outputDataTypes.size(); + this.needPassThrough = properChannelCount != resultDataTypes.size(); this.partitionState = null; - this.properBlockBuilder = new TsBlockBuilder(outputDataTypes.subList(0, properChannelCount)); + this.properBlockBuilder = new TsBlockBuilder(resultDataTypes.subList(0, properChannelCount)); this.maxTsBlockSizeInBytes = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); this.maxTsBlockLineNumber = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); + this.estimatedSerializedBlockFixedSizeInBytes = + getEstimatedSerializedBlockFixedSizeInBytes(resultDataTypes.size()); + this.estimatedSerializedPositionFixedSizeInBytes = + getEstimatedSerializedPositionFixedSizeInBytes(resultDataTypes); + this.variableWidthColumnIndexes = getVariableWidthColumnIndexes(resultDataTypes); this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); - this.tsBlockSerde = new TsBlockSerde(); this.requireRecordSnapshot = requireRecordSnapshot; this.ioTDBLocal = ioTDBLocal; } @@ -216,8 +228,7 @@ private ColumnBuilder getPassThroughIndexBuilder() { } private List buildTsBlock( - List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) - throws IOException { + List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { int positionCount = 0; if (properChannelCount > 0) { // if there is proper column, use its position count @@ -258,41 +269,123 @@ private List buildTsBlock( return result; } - private void addSplitTsBlocks(List 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; + /** + * Splits the final result using a conservative serialized-size upper bound. + * + *

Serializing candidate regions to find their exact sizes would write every value into + * temporary buffers, only for the exchange layer to serialize the selected regions again. + * Rebuilding the result with a size-tracking {@link TsBlockBuilder} would avoid that temporary + * serialization, but it would turn the UDF's batched column output into row-by-row, + * column-by-column copies. + * + *

Instead, fixed-width values are accounted for directly from their data types, while only the + * lengths of variable-width values are inspected. The estimate also reserves space for null + * indicators, dictionary indexes, and encoding metadata, so the currently supported encodings + * cannot make the serialized result exceed the estimate. The resulting regions are views over the + * original columns and do not copy their values. + */ + private void addSplitTsBlocks(List result, TsBlock source) { + if (variableWidthColumnIndexes.length == 0) { + addFixedWidthTsBlocks(result, source); + return; + } + + int regionOffset = 0; + int regionPositionCount = 0; + long estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + Column[] variableWidthColumns = new Column[variableWidthColumnIndexes.length]; + for (int i = 0; i < variableWidthColumnIndexes.length; i++) { + variableWidthColumns[i] = source.getColumn(variableWidthColumnIndexes[i]); + } + for (int position = 0; position < source.getPositionCount(); position++) { + long estimatedPositionSizeInBytes = + getEstimatedPositionSizeInBytes(variableWidthColumns, position); + if (regionPositionCount > 0 + && (regionPositionCount >= maxTsBlockLineNumber + || estimatedRegionSizeInBytes + estimatedPositionSizeInBytes + > maxTsBlockSizeInBytes)) { + result.add(source.getRegion(regionOffset, regionPositionCount)); + regionOffset = position; + regionPositionCount = 0; + estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + } + + estimatedRegionSizeInBytes += estimatedPositionSizeInBytes; + regionPositionCount++; + if (regionPositionCount >= maxTsBlockLineNumber + || estimatedRegionSizeInBytes >= maxTsBlockSizeInBytes) { + result.add(source.getRegion(regionOffset, regionPositionCount)); + regionOffset = position + 1; + regionPositionCount = 0; + estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + } + } + if (regionPositionCount > 0) { + result.add(source.getRegion(regionOffset, regionPositionCount)); } } - private int getMaxSerializedPositionCount(TsBlock source, int offset, int candidatePositionCount) - throws IOException { - if (getSerializedSizeInBytes(source, offset, candidatePositionCount) <= maxTsBlockSizeInBytes) { - return candidatePositionCount; + private void addFixedWidthTsBlocks(List result, TsBlock source) { + long availableSizeInBytes = maxTsBlockSizeInBytes - estimatedSerializedBlockFixedSizeInBytes; + int maxPositionCountBySize = + availableSizeInBytes <= 0 + ? 1 + : (int) + Math.max( + 1, + Math.min( + Integer.MAX_VALUE, + availableSizeInBytes / estimatedSerializedPositionFixedSizeInBytes)); + int maxPositionCount = Math.min(maxTsBlockLineNumber, maxPositionCountBySize); + for (int offset = 0; offset < source.getPositionCount(); offset += maxPositionCount) { + result.add( + source.getRegion(offset, Math.min(maxPositionCount, source.getPositionCount() - offset))); } + } - // 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; + private long getEstimatedPositionSizeInBytes(Column[] variableWidthColumns, int position) { + long sizeInBytes = estimatedSerializedPositionFixedSizeInBytes; + for (Column column : variableWidthColumns) { + if (!column.isNull(position)) { + sizeInBytes += column.getBinary(position).getLength(); } } - return left; + return sizeInBytes; } - private int getSerializedSizeInBytes(TsBlock source, int offset, int positionCount) - throws IOException { - return tsBlockSerde.serialize(source.getRegion(offset, positionCount)).remaining(); + private static long getEstimatedSerializedBlockFixedSizeInBytes(int valueColumnCount) { + int serializedColumnCount = valueColumnCount + 1; + return 2L * Integer.BYTES + + (long) valueColumnCount * TSDataType.getSerializedSize() + + (long) serializedColumnCount + * (Byte.BYTES + ESTIMATED_SERIALIZED_COLUMN_METADATA_SIZE_IN_BYTES); + } + + private static long getEstimatedSerializedPositionFixedSizeInBytes( + List outputDataTypes) { + long sizeInBytes = Long.BYTES + ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES; + for (TSDataType dataType : outputDataTypes) { + sizeInBytes += ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES; + sizeInBytes += dataType.isBinary() ? Integer.BYTES : dataType.getDataTypeSize(); + } + return sizeInBytes; + } + + private static int[] getVariableWidthColumnIndexes(List outputDataTypes) { + int variableWidthColumnCount = 0; + for (TSDataType dataType : outputDataTypes) { + if (dataType.isBinary()) { + variableWidthColumnCount++; + } + } + int[] columnIndexes = new int[variableWidthColumnCount]; + int index = 0; + for (int columnIndex = 0; columnIndex < outputDataTypes.size(); columnIndex++) { + if (outputDataTypes.get(columnIndex).isBinary()) { + columnIndexes[index++] = columnIndex; + } + } + return columnIndexes; } private void consumeCurrentPartitionState() { 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 bc10de7e6abd5..c6f1f837bac82 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 @@ -297,12 +297,12 @@ public void testPartitionRecognizer() { } @Test - public void testVariableWidthResultsAreSplitBySerializedSize() throws Exception { - assertVariableWidthResultsAreSplitBySerializedSize(false); - assertVariableWidthResultsAreSplitBySerializedSize(true); + public void testVariableWidthResultsAreSplitByEstimatedSize() throws Exception { + assertVariableWidthResultsAreSplitByEstimatedSize(false); + assertVariableWidthResultsAreSplitByEstimatedSize(true); } - private void assertVariableWidthResultsAreSplitBySerializedSize(boolean withPassThrough) + private void assertVariableWidthResultsAreSplitByEstimatedSize(boolean withPassThrough) throws Exception { QueryId queryId = new QueryId("large_finish_result_" + withPassThrough); FragmentInstanceId instanceId = @@ -319,7 +319,8 @@ private void assertVariableWidthResultsAreSplitBySerializedSize(boolean withPass int maxBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); Assert.assertTrue(maxLineNumber >= 3); int wideRowCount = maxLineNumber + 1; - int wideValueLength = maxBlockSize / (maxLineNumber - 1) + 1; + int maxWideRowsPerBlock = 8; + int wideValueLength = maxBlockSize / maxWideRowsPerBlock + 1; Binary narrowValue = new Binary("narrow", TSFileConfig.STRING_CHARSET); Binary wideValue = new Binary(new byte[wideValueLength]); TableFunctionProcessorProvider provider = @@ -342,8 +343,8 @@ public void process( @Override public void finish( List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { - // This result exceeds both limits and must be split after pass-through columns are - // appended. + // The result must be sliced into bounded regions after pass-through columns are + // appended, without rebuilding or serializing it merely to determine its size. for (int i = 0; i < wideRowCount; i++) { properColumnBuilders.get(0).writeBinary(wideValue); if (passThroughIndexBuilder != null) { @@ -411,6 +412,7 @@ public long ramBytesUsed() { int returnedRows = 0; int returnedBlocks = 0; + TsBlockSerde serde = new TsBlockSerde(); try (TableFunctionOperator operator = new TableFunctionOperator( operatorContext, @@ -436,7 +438,10 @@ public long ramBytesUsed() { returnedBlocks++; returnedRows += block.getPositionCount(); Assert.assertTrue(block.getPositionCount() <= maxLineNumber); - Assert.assertTrue(new TsBlockSerde().serialize(block).remaining() <= maxBlockSize); + Assert.assertTrue(serde.serialize(block).remaining() <= maxBlockSize); + if (block.getColumn(0).getBinary(0).getLength() == wideValueLength) { + Assert.assertTrue(block.getPositionCount() <= maxWideRowsPerBlock); + } for (int i = 0; i < block.getPositionCount(); i++) { int expectedLength = returnedRows - block.getPositionCount() + i == 0 ? 6 : wideValueLength; @@ -449,7 +454,7 @@ public long ramBytesUsed() { } assertEquals(wideRowCount + 1, returnedRows); - Assert.assertTrue(returnedBlocks > 2); + Assert.assertTrue(returnedBlocks > maxWideRowsPerBlock); } private void checkIteratorSimply(Slice slice, List> expected) { From c8c4b3df293fa93e859866a5cdb82b7b9474d932 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 30 Jul 2026 14:43:43 +0800 Subject: [PATCH 4/6] Estimate table UDF result blocks by memory size --- .../function/TableFunctionOperator.java | 73 +++++++------------ .../tvf/TableFunctionOperatorTest.java | 28 ++++--- 2 files changed, 46 insertions(+), 55 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 289cf63b49423..49988c6a4eb05 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 @@ -40,9 +40,11 @@ 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.BinaryColumn; 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.TimeColumn; import org.apache.tsfile.utils.RamUsageEstimator; import java.util.ArrayList; @@ -61,13 +63,6 @@ public class TableFunctionOperator implements ProcessOperator { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(AggregationMergeSortOperator.class); - // Allow for the null marker, nested encoding metadata, and dictionary metadata of each column. - // These values deliberately overestimate array-encoded columns so the result can be sliced - // without serializing it first. - private static final int ESTIMATED_SERIALIZED_COLUMN_METADATA_SIZE_IN_BYTES = - 2 * Integer.BYTES + 2 * Byte.BYTES; - private static final int ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES = - Byte.BYTES + Integer.BYTES; private final CommonOperatorContext operatorContext; private final Operator inputOperator; @@ -76,8 +71,7 @@ public class TableFunctionOperator implements ProcessOperator { private final TsBlockBuilder properBlockBuilder; private final int maxTsBlockSizeInBytes; private final int maxTsBlockLineNumber; - private final long estimatedSerializedBlockFixedSizeInBytes; - private final long estimatedSerializedPositionFixedSizeInBytes; + private final long estimatedFixedSizePerPositionInBytes; private final int[] variableWidthColumnIndexes; private final int properChannelCount; private final boolean needPassThrough; @@ -123,10 +117,8 @@ public TableFunctionOperator( TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); this.maxTsBlockLineNumber = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); - this.estimatedSerializedBlockFixedSizeInBytes = - getEstimatedSerializedBlockFixedSizeInBytes(resultDataTypes.size()); - this.estimatedSerializedPositionFixedSizeInBytes = - getEstimatedSerializedPositionFixedSizeInBytes(resultDataTypes); + this.estimatedFixedSizePerPositionInBytes = + getEstimatedFixedSizePerPositionInBytes(resultDataTypes); this.variableWidthColumnIndexes = getVariableWidthColumnIndexes(resultDataTypes); this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); @@ -270,7 +262,8 @@ private List buildTsBlock( } /** - * Splits the final result using a conservative serialized-size upper bound. + * Splits the final result using the same logical in-memory size accounting as {@link + * TsBlockBuilder}. * *

Serializing candidate regions to find their exact sizes would write every value into * temporary buffers, only for the exchange layer to serialize the selected regions again. @@ -279,10 +272,10 @@ private List buildTsBlock( * column-by-column copies. * *

Instead, fixed-width values are accounted for directly from their data types, while only the - * lengths of variable-width values are inspected. The estimate also reserves space for null - * indicators, dictionary indexes, and encoding metadata, so the currently supported encodings - * cannot make the serialized result exceed the estimate. The resulting regions are views over the - * original columns and do not copy their values. + * retained sizes of variable-width values are inspected. This deliberately estimates the + * in-memory TsBlock size rather than its serialized size because the two representations are not + * equivalent. The resulting regions are views over the original columns and do not copy their + * values. */ private void addSplitTsBlocks(List result, TsBlock source) { if (variableWidthColumnIndexes.length == 0) { @@ -292,7 +285,7 @@ private void addSplitTsBlocks(List result, TsBlock source) { int regionOffset = 0; int regionPositionCount = 0; - long estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + long estimatedRegionSizeInBytes = 0; Column[] variableWidthColumns = new Column[variableWidthColumnIndexes.length]; for (int i = 0; i < variableWidthColumnIndexes.length; i++) { variableWidthColumns[i] = source.getColumn(variableWidthColumnIndexes[i]); @@ -307,7 +300,7 @@ private void addSplitTsBlocks(List result, TsBlock source) { result.add(source.getRegion(regionOffset, regionPositionCount)); regionOffset = position; regionPositionCount = 0; - estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + estimatedRegionSizeInBytes = 0; } estimatedRegionSizeInBytes += estimatedPositionSizeInBytes; @@ -317,7 +310,7 @@ private void addSplitTsBlocks(List result, TsBlock source) { result.add(source.getRegion(regionOffset, regionPositionCount)); regionOffset = position + 1; regionPositionCount = 0; - estimatedRegionSizeInBytes = estimatedSerializedBlockFixedSizeInBytes; + estimatedRegionSizeInBytes = 0; } } if (regionPositionCount > 0) { @@ -326,16 +319,13 @@ private void addSplitTsBlocks(List result, TsBlock source) { } private void addFixedWidthTsBlocks(List result, TsBlock source) { - long availableSizeInBytes = maxTsBlockSizeInBytes - estimatedSerializedBlockFixedSizeInBytes; int maxPositionCountBySize = - availableSizeInBytes <= 0 - ? 1 - : (int) - Math.max( - 1, - Math.min( - Integer.MAX_VALUE, - availableSizeInBytes / estimatedSerializedPositionFixedSizeInBytes)); + (int) + Math.max( + 1, + Math.min( + Integer.MAX_VALUE, + maxTsBlockSizeInBytes / estimatedFixedSizePerPositionInBytes)); int maxPositionCount = Math.min(maxTsBlockLineNumber, maxPositionCountBySize); for (int offset = 0; offset < source.getPositionCount(); offset += maxPositionCount) { result.add( @@ -344,29 +334,22 @@ private void addFixedWidthTsBlocks(List result, TsBlock source) { } private long getEstimatedPositionSizeInBytes(Column[] variableWidthColumns, int position) { - long sizeInBytes = estimatedSerializedPositionFixedSizeInBytes; + long sizeInBytes = estimatedFixedSizePerPositionInBytes; for (Column column : variableWidthColumns) { if (!column.isNull(position)) { - sizeInBytes += column.getBinary(position).getLength(); + sizeInBytes += column.getBinary(position).ramBytesUsed(); } } return sizeInBytes; } - private static long getEstimatedSerializedBlockFixedSizeInBytes(int valueColumnCount) { - int serializedColumnCount = valueColumnCount + 1; - return 2L * Integer.BYTES - + (long) valueColumnCount * TSDataType.getSerializedSize() - + (long) serializedColumnCount - * (Byte.BYTES + ESTIMATED_SERIALIZED_COLUMN_METADATA_SIZE_IN_BYTES); - } - - private static long getEstimatedSerializedPositionFixedSizeInBytes( - List outputDataTypes) { - long sizeInBytes = Long.BYTES + ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES; + private static long getEstimatedFixedSizePerPositionInBytes(List outputDataTypes) { + long sizeInBytes = TimeColumn.SIZE_IN_BYTES_PER_POSITION; for (TSDataType dataType : outputDataTypes) { - sizeInBytes += ESTIMATED_SERIALIZED_POSITION_OVERHEAD_IN_BYTES; - sizeInBytes += dataType.isBinary() ? Integer.BYTES : dataType.getDataTypeSize(); + sizeInBytes += + dataType.isBinary() + ? BinaryColumn.SHALLOW_SIZE_IN_BYTES_PER_POSITION + : dataType.getDataTypeSize() + Byte.BYTES; } return sizeInBytes; } 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 c6f1f837bac82..e1ea062478002 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 @@ -44,8 +44,10 @@ 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.BinaryColumn; +import org.apache.tsfile.read.common.block.column.LongColumn; import org.apache.tsfile.read.common.block.column.RunLengthEncodedColumn; -import org.apache.tsfile.read.common.block.column.TsBlockSerde; +import org.apache.tsfile.read.common.block.column.TimeColumn; import org.apache.tsfile.utils.Binary; import org.junit.AfterClass; import org.junit.Assert; @@ -297,12 +299,12 @@ public void testPartitionRecognizer() { } @Test - public void testVariableWidthResultsAreSplitByEstimatedSize() throws Exception { - assertVariableWidthResultsAreSplitByEstimatedSize(false); - assertVariableWidthResultsAreSplitByEstimatedSize(true); + public void testVariableWidthResultsAreSplitByEstimatedMemorySize() throws Exception { + assertVariableWidthResultsAreSplitByEstimatedMemorySize(false); + assertVariableWidthResultsAreSplitByEstimatedMemorySize(true); } - private void assertVariableWidthResultsAreSplitByEstimatedSize(boolean withPassThrough) + private void assertVariableWidthResultsAreSplitByEstimatedMemorySize(boolean withPassThrough) throws Exception { QueryId queryId = new QueryId("large_finish_result_" + withPassThrough); FragmentInstanceId instanceId = @@ -343,8 +345,8 @@ public void process( @Override public void finish( List properColumnBuilders, ColumnBuilder passThroughIndexBuilder) { - // The result must be sliced into bounded regions after pass-through columns are - // appended, without rebuilding or serializing it merely to determine its size. + // The result must be sliced according to its estimated in-memory size after + // pass-through columns are appended, without rebuilding or serializing it. for (int i = 0; i < wideRowCount; i++) { properColumnBuilders.get(0).writeBinary(wideValue); if (passThroughIndexBuilder != null) { @@ -412,7 +414,6 @@ public long ramBytesUsed() { int returnedRows = 0; int returnedBlocks = 0; - TsBlockSerde serde = new TsBlockSerde(); try (TableFunctionOperator operator = new TableFunctionOperator( operatorContext, @@ -438,18 +439,25 @@ public long ramBytesUsed() { returnedBlocks++; returnedRows += block.getPositionCount(); Assert.assertTrue(block.getPositionCount() <= maxLineNumber); - Assert.assertTrue(serde.serialize(block).remaining() <= maxBlockSize); if (block.getColumn(0).getBinary(0).getLength() == wideValueLength) { Assert.assertTrue(block.getPositionCount() <= maxWideRowsPerBlock); } + long estimatedBlockSize = 0; for (int i = 0; i < block.getPositionCount(); i++) { + Binary value = block.getColumn(0).getBinary(i); + estimatedBlockSize += + TimeColumn.SIZE_IN_BYTES_PER_POSITION + + BinaryColumn.SHALLOW_SIZE_IN_BYTES_PER_POSITION + + value.ramBytesUsed() + + (withPassThrough ? LongColumn.SIZE_IN_BYTES_PER_POSITION : 0); int expectedLength = returnedRows - block.getPositionCount() + i == 0 ? 6 : wideValueLength; - assertEquals(expectedLength, block.getColumn(0).getBinary(i).getLength()); + assertEquals(expectedLength, value.getLength()); if (withPassThrough) { assertEquals(1, block.getColumn(1).getLong(i)); } } + Assert.assertTrue(estimatedBlockSize <= maxBlockSize); } } From 356c516b0f4ac6dce50335634f42406bdd20ae3e Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 30 Jul 2026 15:35:13 +0800 Subject: [PATCH 5/6] Clarify table UDF result size accounting --- .../function/TableFunctionOperator.java | 27 +++++--- .../tvf/TableFunctionOperatorTest.java | 64 +++++++++++++++++-- 2 files changed, 76 insertions(+), 15 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 49988c6a4eb05..90b530402fbbe 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 @@ -71,7 +71,7 @@ public class TableFunctionOperator implements ProcessOperator { private final TsBlockBuilder properBlockBuilder; private final int maxTsBlockSizeInBytes; private final int maxTsBlockLineNumber; - private final long estimatedFixedSizePerPositionInBytes; + private final long estimatedBaseSizePerPositionInBytes; private final int[] variableWidthColumnIndexes; private final int properChannelCount; private final boolean needPassThrough; @@ -117,8 +117,8 @@ public TableFunctionOperator( TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); this.maxTsBlockLineNumber = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); - this.estimatedFixedSizePerPositionInBytes = - getEstimatedFixedSizePerPositionInBytes(resultDataTypes); + this.estimatedBaseSizePerPositionInBytes = + getEstimatedBaseSizePerPositionInBytes(resultDataTypes); this.variableWidthColumnIndexes = getVariableWidthColumnIndexes(resultDataTypes); this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); @@ -276,10 +276,20 @@ private List buildTsBlock( * in-memory TsBlock size rather than its serialized size because the two representations are not * equivalent. The resulting regions are views over the original columns and do not copy their * values. + * + *

For a block containing only fixed-width columns, the number of positions in each region is + * calculated directly, so no per-position scan is required. If the block contains variable-width + * columns, every position in those columns must be inspected once because their payload sizes can + * differ by row. This costs {@code O(positionCount * variableWidthColumnCount)}, but each + * inspection only reads the null state, the {@link org.apache.tsfile.utils.Binary} reference, and + * its allocated size. It does not traverse the payload bytes, serialize values, or copy column + * data. A shared {@code Binary} referenced by multiple positions is deliberately counted once per + * position, which may produce smaller regions but keeps the estimate conservative. */ private void addSplitTsBlocks(List result, TsBlock source) { if (variableWidthColumnIndexes.length == 0) { - addFixedWidthTsBlocks(result, source); + // This fast path is used only when every column has a fixed per-position size. + addFixedWidthOnlyTsBlocks(result, source); return; } @@ -318,14 +328,14 @@ private void addSplitTsBlocks(List result, TsBlock source) { } } - private void addFixedWidthTsBlocks(List result, TsBlock source) { + private void addFixedWidthOnlyTsBlocks(List result, TsBlock source) { int maxPositionCountBySize = (int) Math.max( 1, Math.min( Integer.MAX_VALUE, - maxTsBlockSizeInBytes / estimatedFixedSizePerPositionInBytes)); + maxTsBlockSizeInBytes / estimatedBaseSizePerPositionInBytes)); int maxPositionCount = Math.min(maxTsBlockLineNumber, maxPositionCountBySize); for (int offset = 0; offset < source.getPositionCount(); offset += maxPositionCount) { result.add( @@ -334,7 +344,8 @@ private void addFixedWidthTsBlocks(List result, TsBlock source) { } private long getEstimatedPositionSizeInBytes(Column[] variableWidthColumns, int position) { - long sizeInBytes = estimatedFixedSizePerPositionInBytes; + // Fixed-width values and variable-width payloads consume the same block-size budget. + long sizeInBytes = estimatedBaseSizePerPositionInBytes; for (Column column : variableWidthColumns) { if (!column.isNull(position)) { sizeInBytes += column.getBinary(position).ramBytesUsed(); @@ -343,7 +354,7 @@ private long getEstimatedPositionSizeInBytes(Column[] variableWidthColumns, int return sizeInBytes; } - private static long getEstimatedFixedSizePerPositionInBytes(List outputDataTypes) { + private static long getEstimatedBaseSizePerPositionInBytes(List outputDataTypes) { long sizeInBytes = TimeColumn.SIZE_IN_BYTES_PER_POSITION; for (TSDataType dataType : outputDataTypes) { sizeInBytes += 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 e1ea062478002..bfc4541c5a1b5 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 @@ -304,9 +304,49 @@ public void testVariableWidthResultsAreSplitByEstimatedMemorySize() throws Excep assertVariableWidthResultsAreSplitByEstimatedMemorySize(true); } + @Test + public void testFixedAndVariableWidthColumnsShareSizeBudget() throws Exception { + int originalMaxBlockSize = + TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + Binary wideValue = new Binary(new byte[64]); + long variableWidthSizePerPosition = + BinaryColumn.SHALLOW_SIZE_IN_BYTES_PER_POSITION + wideValue.ramBytesUsed(); + long fixedWidthSizePerPosition = + TimeColumn.SIZE_IN_BYTES_PER_POSITION + LongColumn.SIZE_IN_BYTES_PER_POSITION; + int combinedMaxBlockSize = + Math.toIntExact(2 * Math.max(variableWidthSizePerPosition, fixedWidthSizePerPosition)); + + // Two rows of either part fit independently, but two complete rows do not. + Assert.assertTrue(2 * variableWidthSizePerPosition <= combinedMaxBlockSize); + Assert.assertTrue(2 * fixedWidthSizePerPosition <= combinedMaxBlockSize); + Assert.assertTrue( + 2 * (variableWidthSizePerPosition + fixedWidthSizePerPosition) > combinedMaxBlockSize); + try { + assertVariableWidthResultsAreSplitByEstimatedMemorySize( + true, 3, wideValue, 1, combinedMaxBlockSize); + } finally { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(originalMaxBlockSize); + } + } + private void assertVariableWidthResultsAreSplitByEstimatedMemorySize(boolean withPassThrough) throws Exception { - QueryId queryId = new QueryId("large_finish_result_" + withPassThrough); + int maxLineNumber = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockLineNumber(); + int maxBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + int maxWideRowsPerBlock = 8; + Binary wideValue = new Binary(new byte[maxBlockSize / maxWideRowsPerBlock + 1]); + assertVariableWidthResultsAreSplitByEstimatedMemorySize( + withPassThrough, maxLineNumber + 1, wideValue, maxWideRowsPerBlock, null); + } + + private void assertVariableWidthResultsAreSplitByEstimatedMemorySize( + boolean withPassThrough, + int wideRowCount, + Binary wideValue, + int maxWideRowsPerBlock, + Integer maxBlockSizeOverride) + throws Exception { + QueryId queryId = new QueryId("large_finish_result_" + withPassThrough + "_" + wideRowCount); FragmentInstanceId instanceId = new FragmentInstanceId(new PlanFragmentId(queryId, 0), "stub-instance"); FragmentInstanceStateMachine stateMachine = @@ -318,13 +358,9 @@ private void assertVariableWidthResultsAreSplitByEstimatedMemorySize(boolean wit 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 maxWideRowsPerBlock = 8; - int wideValueLength = maxBlockSize / maxWideRowsPerBlock + 1; + int wideValueLength = wideValue.getLength(); Binary narrowValue = new Binary("narrow", TSFileConfig.STRING_CHARSET); - Binary wideValue = new Binary(new byte[wideValueLength]); TableFunctionProcessorProvider provider = new TableFunctionProcessorProvider() { @Override @@ -414,6 +450,13 @@ public long ramBytesUsed() { int returnedRows = 0; int returnedBlocks = 0; + if (maxBlockSizeOverride != null) { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(maxBlockSizeOverride); + } + int maxBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + if (maxBlockSizeOverride != null) { + assertEquals(maxBlockSizeOverride.intValue(), maxBlockSize); + } try (TableFunctionOperator operator = new TableFunctionOperator( operatorContext, @@ -440,7 +483,14 @@ public long ramBytesUsed() { returnedRows += block.getPositionCount(); Assert.assertTrue(block.getPositionCount() <= maxLineNumber); if (block.getColumn(0).getBinary(0).getLength() == wideValueLength) { - Assert.assertTrue(block.getPositionCount() <= maxWideRowsPerBlock); + Assert.assertTrue( + "wide block position count: " + + block.getPositionCount() + + ", expected at most: " + + maxWideRowsPerBlock + + ", configured size: " + + maxBlockSize, + block.getPositionCount() <= maxWideRowsPerBlock); } long estimatedBlockSize = 0; for (int i = 0; i < block.getPositionCount(); i++) { From c380a825a1a53026d008fd99234d1f81f3520686 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 30 Jul 2026 16:01:48 +0800 Subject: [PATCH 6/6] Remove redundant table UDF split condition --- .../operator/process/function/TableFunctionOperator.java | 4 +--- 1 file changed, 1 insertion(+), 3 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 90b530402fbbe..0e371884809e4 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 @@ -304,9 +304,7 @@ private void addSplitTsBlocks(List result, TsBlock source) { long estimatedPositionSizeInBytes = getEstimatedPositionSizeInBytes(variableWidthColumns, position); if (regionPositionCount > 0 - && (regionPositionCount >= maxTsBlockLineNumber - || estimatedRegionSizeInBytes + estimatedPositionSizeInBytes - > maxTsBlockSizeInBytes)) { + && estimatedRegionSizeInBytes + estimatedPositionSizeInBytes > maxTsBlockSizeInBytes) { result.add(source.getRegion(regionOffset, regionPositionCount)); regionOffset = position; regionPositionCount = 0;