diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java new file mode 100644 index 0000000000000..4a032af59016c --- /dev/null +++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/LargeResultTableFunction.java @@ -0,0 +1,136 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.query.udf.example.relational; + +import org.apache.iotdb.udf.api.exception.UDFException; +import org.apache.iotdb.udf.api.relational.TableFunction; +import org.apache.iotdb.udf.api.relational.access.Record; +import org.apache.iotdb.udf.api.relational.table.MapTableFunctionHandle; +import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis; +import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle; +import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider; +import org.apache.iotdb.udf.api.relational.table.argument.Argument; +import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema; +import org.apache.iotdb.udf.api.relational.table.argument.ScalarArgument; +import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor; +import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification; +import org.apache.iotdb.udf.api.relational.table.specification.ScalarParameterSpecification; +import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification; +import org.apache.iotdb.udf.api.type.Type; + +import org.apache.tsfile.block.column.ColumnBuilder; +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.utils.Binary; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +public class LargeResultTableFunction implements TableFunction { + + private static final String TABLE_PARAMETER_NAME = "DATA"; + private static final String REPEAT_COUNT_PARAMETER_NAME = "REPEAT_COUNT"; + private static final String PAYLOAD_SIZE_PARAMETER_NAME = "PAYLOAD_SIZE"; + + @Override + public List getArgumentsSpecifications() { + return Arrays.asList( + TableParameterSpecification.builder() + .name(TABLE_PARAMETER_NAME) + .rowSemantics() + .passThroughColumns() + .build(), + ScalarParameterSpecification.builder() + .name(REPEAT_COUNT_PARAMETER_NAME) + .type(Type.INT32) + .build(), + ScalarParameterSpecification.builder() + .name(PAYLOAD_SIZE_PARAMETER_NAME) + .type(Type.INT32) + .build()); + } + + @Override + public TableFunctionAnalysis analyze(Map arguments) throws UDFException { + MapTableFunctionHandle handle = + new MapTableFunctionHandle.Builder() + .addProperty( + REPEAT_COUNT_PARAMETER_NAME, + ((ScalarArgument) arguments.get(REPEAT_COUNT_PARAMETER_NAME)).getValue()) + .addProperty( + PAYLOAD_SIZE_PARAMETER_NAME, + ((ScalarArgument) arguments.get(PAYLOAD_SIZE_PARAMETER_NAME)).getValue()) + .build(); + return TableFunctionAnalysis.builder() + .properColumnSchema( + DescribedSchema.builder() + .addField("repeat_index", Type.INT32) + .addField("payload", Type.STRING) + .build()) + .requiredColumns(TABLE_PARAMETER_NAME, Collections.singletonList(0)) + .handle(handle) + .build(); + } + + @Override + public TableFunctionHandle createTableFunctionHandle() { + return new MapTableFunctionHandle(); + } + + @Override + public TableFunctionProcessorProvider getProcessorProvider( + TableFunctionHandle tableFunctionHandle) { + return new TableFunctionProcessorProvider() { + @Override + public TableFunctionDataProcessor getDataProcessor() { + return new TableFunctionDataProcessor() { + private final int repeatCount = + (int) + ((MapTableFunctionHandle) tableFunctionHandle) + .getProperty(REPEAT_COUNT_PARAMETER_NAME); + private final String payloadSuffix = + "x" + .repeat( + (int) + ((MapTableFunctionHandle) tableFunctionHandle) + .getProperty(PAYLOAD_SIZE_PARAMETER_NAME)); + private long recordIndex; + + @Override + public void process( + Record input, + List properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + for (int repeatIndex = 0; repeatIndex < repeatCount; repeatIndex++) { + properColumnBuilders.get(0).writeInt(repeatIndex); + properColumnBuilders + .get(1) + .writeBinary( + new Binary(repeatIndex + ":" + payloadSuffix, TSFileConfig.STRING_CHARSET)); + passThroughIndexBuilder.writeLong(recordIndex); + } + recordIndex++; + } + }; + } + }; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java index 7fe6648fe70ba..9db3acb13e098 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java @@ -26,13 +26,18 @@ import org.junit.After; import org.junit.AfterClass; +import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; +import java.util.HashSet; +import java.util.Set; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; @@ -42,6 +47,9 @@ @Category({TableLocalStandaloneIT.class, TableClusterIT.class}) public class IoTDBUserDefinedTableFunctionIT { private static final String DATABASE_NAME = "test"; + private static final int MAX_TSBLOCK_SIZE_IN_BYTES = 1024; + private static final int LARGE_RESULT_REPEAT_COUNT = 64; + private static final int LARGE_RESULT_PAYLOAD_SIZE = 128; private static final String[] sqls = new String[] { "CREATE DATABASE " + DATABASE_NAME, @@ -57,6 +65,10 @@ public class IoTDBUserDefinedTableFunctionIT { @BeforeClass public static void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getDataNodeCommonConfig() + .setMaxTsBlockSizeInByte(MAX_TSBLOCK_SIZE_IN_BYTES); EnvFactory.getEnv().initClusterEnvironment(); insertData(); } @@ -182,6 +194,74 @@ public void testMyRepeat() { DATABASE_NAME); } + @Test + public void testLargeResultIsSplitWithoutDataLoss() throws Exception { + SQLFunctionUtils.createUDF( + "large_result", + "org.apache.iotdb.db.query.udf.example.relational.LargeResultTableFunction"); + + Set returnedRows = new HashSet<>(); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE_NAME); + try (ResultSet resultSet = + statement.executeQuery( + "SELECT * FROM large_result(vehicle, " + + LARGE_RESULT_REPEAT_COUNT + + ", " + + LARGE_RESULT_PAYLOAD_SIZE + + ")")) { + while (resultSet.next()) { + int repeatIndex = resultSet.getInt("repeat_index"); + long time = resultSet.getLong("time"); + Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT); + Assert.assertEquals( + repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE), + resultSet.getString("payload")); + assertPassThroughColumns(resultSet, time); + Assert.assertTrue( + "Duplicate result row for time " + time + " and repeat index " + repeatIndex, + returnedRows.add(time + ":" + repeatIndex)); + } + } + } + + long[] inputTimes = new long[] {1, 2, 3, 5}; + Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size()); + for (long time : inputTimes) { + for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) { + Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex)); + } + } + } + + private static void assertPassThroughColumns(ResultSet resultSet, long time) throws SQLException { + switch ((int) time) { + case 1: + Assert.assertEquals("d0", resultSet.getString("device_id")); + Assert.assertEquals(1, resultSet.getInt("s1")); + Assert.assertEquals(1, resultSet.getLong("s2")); + break; + case 2: + Assert.assertEquals("d0", resultSet.getString("device_id")); + Assert.assertNull(resultSet.getObject("s1")); + Assert.assertEquals(2, resultSet.getLong("s2")); + break; + case 3: + Assert.assertEquals("d0", resultSet.getString("device_id")); + Assert.assertEquals(3, resultSet.getInt("s1")); + Assert.assertEquals(3, resultSet.getLong("s2")); + break; + case 5: + Assert.assertEquals("d1", resultSet.getString("device_id")); + Assert.assertEquals(4, resultSet.getInt("s1")); + Assert.assertNull(resultSet.getObject("s2")); + break; + default: + Assert.fail("Unexpected pass-through time: " + time); + } + } + @Test public void testHybrid() { SQLFunctionUtils.createUDF( diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java index d4bc8a0da4c79..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 @@ -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; @@ -62,14 +64,15 @@ 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 long estimatedBaseSizePerPositionInBytes; + private final int[] variableWidthColumnIndexes; private final int properChannelCount; private final boolean needPassThrough; private final PartitionCache partitionCache; @@ -105,10 +108,18 @@ 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.estimatedBaseSizePerPositionInBytes = + getEstimatedBaseSizePerPositionInBytes(resultDataTypes); + this.variableWidthColumnIndexes = getVariableWidthColumnIndexes(resultDataTypes); this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); this.requireRecordSnapshot = requireRecordSnapshot; @@ -241,16 +252,134 @@ 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; } + /** + * 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. + * 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 + * 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. + * + *

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) { + // This fast path is used only when every column has a fixed per-position size. + addFixedWidthOnlyTsBlocks(result, source); + return; + } + + int regionOffset = 0; + int regionPositionCount = 0; + long estimatedRegionSizeInBytes = 0; + 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 + && estimatedRegionSizeInBytes + estimatedPositionSizeInBytes > maxTsBlockSizeInBytes) { + result.add(source.getRegion(regionOffset, regionPositionCount)); + regionOffset = position; + regionPositionCount = 0; + estimatedRegionSizeInBytes = 0; + } + + estimatedRegionSizeInBytes += estimatedPositionSizeInBytes; + regionPositionCount++; + if (regionPositionCount >= maxTsBlockLineNumber + || estimatedRegionSizeInBytes >= maxTsBlockSizeInBytes) { + result.add(source.getRegion(regionOffset, regionPositionCount)); + regionOffset = position + 1; + regionPositionCount = 0; + estimatedRegionSizeInBytes = 0; + } + } + if (regionPositionCount > 0) { + result.add(source.getRegion(regionOffset, regionPositionCount)); + } + } + + private void addFixedWidthOnlyTsBlocks(List result, TsBlock source) { + int maxPositionCountBySize = + (int) + Math.max( + 1, + Math.min( + Integer.MAX_VALUE, + maxTsBlockSizeInBytes / estimatedBaseSizePerPositionInBytes)); + 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))); + } + } + + private long getEstimatedPositionSizeInBytes(Column[] variableWidthColumns, int position) { + // 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(); + } + } + return sizeInBytes; + } + + private static long getEstimatedBaseSizePerPositionInBytes(List outputDataTypes) { + long sizeInBytes = TimeColumn.SIZE_IN_BYTES_PER_POSITION; + for (TSDataType dataType : outputDataTypes) { + sizeInBytes += + dataType.isBinary() + ? BinaryColumn.SHALLOW_SIZE_IN_BYTES_PER_POSITION + : dataType.getDataTypeSize() + Byte.BYTES; + } + 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() { partitionState = null; } @@ -272,6 +401,7 @@ public boolean hasNext() throws Exception { public void close() throws Exception { partitionCache.close(); inputOperator.close(); + resultTsBlocks.clear(); if (processor != null) { destroyProcessor(processor); processor = null; @@ -281,18 +411,18 @@ public void close() throws Exception { @Override public boolean isFinished() throws Exception { - return finished; + 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 Math.max(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES, properBlockBuilder.getRetainedSizeInBytes()); + 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 bcd17aa665c47..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 @@ -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,21 @@ 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.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.TimeColumn; import org.apache.tsfile.utils.Binary; import org.junit.AfterClass; import org.junit.Assert; @@ -53,6 +63,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 +298,223 @@ public void testPartitionRecognizer() { } } + @Test + public void testVariableWidthResultsAreSplitByEstimatedMemorySize() throws Exception { + assertVariableWidthResultsAreSplitByEstimatedMemorySize(false); + 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 { + 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 = + 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(); + Assert.assertTrue(maxLineNumber >= 3); + int wideValueLength = wideValue.getLength(); + Binary narrowValue = new Binary("narrow", TSFileConfig.STRING_CHARSET); + TableFunctionProcessorProvider provider = + new TableFunctionProcessorProvider() { + @Override + public TableFunctionDataProcessor getDataProcessor() { + return new TableFunctionDataProcessor() { + @Override + public void process( + Record input, + List properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + // 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) { + // 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) { + passThroughIndexBuilder.writeLong(0); + } + } + } + }; + } + }; + + 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; + 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, + provider, + singleRowChild, + Collections.singletonList(TSDataType.INT64), + withPassThrough + ? Arrays.asList(TSDataType.TEXT, TSDataType.INT64) + : Collections.singletonList(TSDataType.TEXT), + 1, + Collections.singletonList(0), + withPassThrough ? Collections.singletonList(0) : Collections.emptyList(), + withPassThrough, + 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); + if (block.getColumn(0).getBinary(0).getLength() == wideValueLength) { + 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++) { + 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, value.getLength()); + if (withPassThrough) { + assertEquals(1, block.getColumn(1).getLong(i)); + } + } + Assert.assertTrue(estimatedBlockSize <= maxBlockSize); + } + } + + assertEquals(wideRowCount + 1, returnedRows); + Assert.assertTrue(returnedBlocks > maxWideRowsPerBlock); + } + private void checkIteratorSimply(Slice slice, List> expected) { Iterator recordIterable = slice.getRequiredRecordIterator(false); int i = 0;