diff --git a/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java index 36a2ec1292e17..37fa74b9bd5a6 100644 --- a/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java +++ b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java @@ -26,6 +26,7 @@ import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; import org.apache.iotdb.rest.protocol.model.ExecutionStatus; import org.apache.iotdb.rest.protocol.v1.model.QueryDataSet; import org.apache.iotdb.rpc.TSStatusCode; @@ -88,6 +89,7 @@ private static final class TestQueryExecution implements IQueryExecution { private final Queue> batches = new ArrayDeque<>(); private final DatasetHeader datasetHeader; + private final PendingResultFragment pendingResultFragment = new PendingResultFragment(() -> {}); private TestQueryExecution(TsBlock... tsBlocks) { for (TsBlock tsBlock : tsBlocks) { @@ -124,10 +126,15 @@ public Optional getBatchResult() throws IoTDBException { } @Override - public Optional getByteBufferBatchResult() { + public Optional getByteBufferBatchResult(boolean retainMemory) { return Optional.empty(); } + @Override + public PendingResultFragment getPendingResultFragment() { + return pendingResultFragment; + } + @Override public boolean hasNextResult() { return !batches.isEmpty(); diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionDataSet.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionDataSet.java index 5076afd52d05b..511805e5b8fd8 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionDataSet.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionDataSet.java @@ -61,7 +61,10 @@ public SessionDataSet( ZoneId zoneId, int timeFactor, boolean tableModel, - List columnIndex2TsBlockColumnIndexList) { + List columnIndex2TsBlockColumnIndexList, + boolean offsetSet, + int offset, + int totalLength) { this.ioTDBRpcDataSet = new IoTDBRpcDataSet( sql, @@ -81,7 +84,48 @@ public SessionDataSet( RpcUtils.DEFAULT_TIME_FORMAT, timeFactor, tableModel, - columnIndex2TsBlockColumnIndexList); + columnIndex2TsBlockColumnIndexList, + offsetSet, + offset, + totalLength); + } + + @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters + public SessionDataSet( + String sql, + List columnNameList, + List columnTypeList, + Map columnNameIndex, + long queryId, + long statementId, + IClientRPCService.Iface client, + long sessionId, + List queryResult, + boolean ignoreTimeStamp, + boolean moreData, + ZoneId zoneId, + int timeFactor, + boolean tableModel, + List columnIndex2TsBlockColumnIndexList) { + this( + sql, + columnNameList, + columnTypeList, + columnNameIndex, + queryId, + statementId, + client, + sessionId, + queryResult, + ignoreTimeStamp, + moreData, + zoneId, + timeFactor, + tableModel, + columnIndex2TsBlockColumnIndexList, + false, + 0, + 0); } @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters @@ -102,7 +146,10 @@ public SessionDataSet( ZoneId zoneId, int timeFactor, boolean tableModel, - List columnIndex2TsBlockColumnIndexList) { + List columnIndex2TsBlockColumnIndexList, + boolean offsetSet, + int offset, + int totalLength) { this.ioTDBRpcDataSet = new IoTDBRpcDataSet( sql, @@ -122,7 +169,52 @@ public SessionDataSet( RpcUtils.DEFAULT_TIME_FORMAT, timeFactor, tableModel, - columnIndex2TsBlockColumnIndexList); + columnIndex2TsBlockColumnIndexList, + offsetSet, + offset, + totalLength); + } + + @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters + public SessionDataSet( + String sql, + List columnNameList, + List columnTypeList, + Map columnNameIndex, + long queryId, + long statementId, + IClientRPCService.Iface client, + long sessionId, + List queryResult, + boolean ignoreTimeStamp, + long timeout, + boolean moreData, + int fetchSize, + ZoneId zoneId, + int timeFactor, + boolean tableModel, + List columnIndex2TsBlockColumnIndexList) { + this( + sql, + columnNameList, + columnTypeList, + columnNameIndex, + queryId, + statementId, + client, + sessionId, + queryResult, + ignoreTimeStamp, + timeout, + moreData, + fetchSize, + zoneId, + timeFactor, + tableModel, + columnIndex2TsBlockColumnIndexList, + false, + 0, + 0); } public int getFetchSize() { diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Config.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Config.java index 54aca784d15a5..79517b5b938c1 100644 --- a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Config.java +++ b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Config.java @@ -93,4 +93,6 @@ private Config() { public static final String SQL_DIALECT = "sql_dialect"; public static final String DATABASE = "db"; + + public static final String ENABLE_RESULT_CHUNKING = "enable_result_chunking"; } diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java index 96beecffe1fc4..f6385249e6b92 100644 --- a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java +++ b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java @@ -601,6 +601,8 @@ private void openSession() throws SQLException { openReq.setZoneId(getTimeZone()); openReq.putToConfiguration(Config.VERSION, params.getVersion().toString()); openReq.putToConfiguration(Config.SQL_DIALECT, params.getSqlDialect()); + // this client can assemble fragmented query results, see TSFetchResultsReq#offset + openReq.putToConfiguration(Config.ENABLE_RESULT_CHUNKING, Boolean.TRUE.toString()); params.getDb().ifPresent(db -> openReq.putToConfiguration(Config.DATABASE, db)); TSOpenSessionResp openResp = null; diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBJDBCResultSet.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBJDBCResultSet.java index 5d5cffdf5452e..8cf3312be12a2 100644 --- a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBJDBCResultSet.java +++ b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBJDBCResultSet.java @@ -104,7 +104,10 @@ public IoTDBJDBCResultSet( ZoneId zoneId, Charset charset, boolean tableModel, - List columnIndex2TsBlockColumnIndexList) + List columnIndex2TsBlockColumnIndexList, + boolean offsetSet, + int offset, + int totalLength) throws SQLException { this.ioTDBRpcDataSet = new IoTDBRpcDataSet( @@ -125,7 +128,10 @@ public IoTDBJDBCResultSet( timeFormat, statement.getTimeFactor(), tableModel, - columnIndex2TsBlockColumnIndexList); + columnIndex2TsBlockColumnIndexList, + offsetSet, + offset, + totalLength); this.statement = statement; this.columnTypeList = columnTypeList; if (tracingInfo != null) { @@ -135,6 +141,49 @@ public IoTDBJDBCResultSet( this.charset = charset; } + @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters + public IoTDBJDBCResultSet( + IoTDBStatement statement, + List columnNameList, + List columnTypeList, + Map columnNameIndex, + boolean ignoreTimeStamp, + IClientRPCService.Iface client, + String sql, + long queryId, + long sessionId, + List dataSet, + TSTracingInfo tracingInfo, + long timeout, + boolean moreData, + ZoneId zoneId, + Charset charset, + boolean tableModel, + List columnIndex2TsBlockColumnIndexList) + throws SQLException { + this( + statement, + columnNameList, + columnTypeList, + columnNameIndex, + ignoreTimeStamp, + client, + sql, + queryId, + sessionId, + dataSet, + tracingInfo, + timeout, + moreData, + zoneId, + charset, + tableModel, + columnIndex2TsBlockColumnIndexList, + false, + 0, + 0); + } + @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters public IoTDBJDBCResultSet( Statement statement, @@ -150,7 +199,10 @@ public IoTDBJDBCResultSet( TSTracingInfo tracingInfo, long timeout, boolean moreData, - ZoneId zoneId) + ZoneId zoneId, + boolean offsetSet, + int offset, + int totalLength) throws SQLException { this.ioTDBRpcDataSet = new IoTDBRpcDataSet( @@ -171,7 +223,10 @@ public IoTDBJDBCResultSet( timeFormat, ((IoTDBStatement) statement).getTimeFactor(), false, - IntStream.range(0, columnNameList.size()).boxed().collect(Collectors.toList())); + IntStream.range(0, columnNameList.size()).boxed().collect(Collectors.toList()), + offsetSet, + offset, + totalLength); this.statement = (IoTDBStatement) statement; this.columnTypeList = columnTypeList; if (tracingInfo != null) { @@ -180,6 +235,43 @@ public IoTDBJDBCResultSet( } } + @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters + public IoTDBJDBCResultSet( + Statement statement, + List columnNameList, + List columnTypeList, + Map columnNameIndex, + boolean ignoreTimeStamp, + IClientRPCService.Iface client, + String sql, + long queryId, + long sessionId, + List dataSet, + TSTracingInfo tracingInfo, + long timeout, + boolean moreData, + ZoneId zoneId) + throws SQLException { + this( + statement, + columnNameList, + columnTypeList, + columnNameIndex, + ignoreTimeStamp, + client, + sql, + queryId, + sessionId, + dataSet, + tracingInfo, + timeout, + moreData, + zoneId, + false, + 0, + 0); + } + @Override public boolean isWrapperFor(Class iface) throws SQLException { throw new SQLException(Constant.METHOD_NOT_SUPPORTED); diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java index ff790cbd9f438..858462d6b2460 100644 --- a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java +++ b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java @@ -412,7 +412,10 @@ private boolean executeSQL(String sql) throws TException, SQLException { zoneId, charset, execResp.isSetTableModel() && execResp.isTableModel(), - execResp.getColumnIndex2TsBlockColumnIndexList()); + execResp.getColumnIndex2TsBlockColumnIndexList(), + execResp.isSetOffset(), + execResp.offset, + execResp.totalLength); } return true; } @@ -549,7 +552,10 @@ private ResultSet executeQuerySQL(String sql, long timeoutInMS) throws TExceptio zoneId, charset, execResp.isSetTableModel() && execResp.isTableModel(), - execResp.getColumnIndex2TsBlockColumnIndexList()); + execResp.getColumnIndex2TsBlockColumnIndexList(), + execResp.isSetOffset(), + execResp.offset, + execResp.totalLength); } return resultSet; } diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatement.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatement.java index fe4b1a0c1f3d1..22493290fbe32 100644 --- a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatement.java +++ b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBTablePreparedStatement.java @@ -238,7 +238,10 @@ private ResultSet processQueryResult(TSExecuteStatementResp resp) throws SQLExce resp.isSetTableModel() && resp.isTableModel(), resp.getColumnIndex2TsBlockColumnIndexList() == null ? null - : new ArrayList<>(resp.getColumnIndex2TsBlockColumnIndexList())); + : new ArrayList<>(resp.getColumnIndex2TsBlockColumnIndexList()), + resp.isSetOffset(), + resp.offset, + resp.totalLength); return resultSet; } return null; diff --git a/iotdb-client/service-rpc/pom.xml b/iotdb-client/service-rpc/pom.xml index 206cc949fe440..d53db34b388bb 100644 --- a/iotdb-client/service-rpc/pom.xml +++ b/iotdb-client/service-rpc/pom.xml @@ -75,6 +75,11 @@ junit test + + org.mockito + mockito-core + test + org.slf4j slf4j-api diff --git a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java index cd925b4863af1..862e15b3b90e8 100644 --- a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -127,4 +127,15 @@ private RpcMessages() {} public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 = "Data type %s is not supported."; + + public static final String EXCEPTION_CHUNKED_RESULT_EMPTY_FRAGMENT_03A0EF39 = + "Empty fragment received when assembling a chunked query result"; + public static final String + EXCEPTION_CHUNKED_RESULT_FRAGMENT_OFFSET_ARG_NOT_GREATER_THAN_ACCUMULATED_LENGTH_ARG_A0F032E3 = + "Fragment offset %d is not greater than the accumulated length %d"; + public static final String EXCEPTION_CHUNKED_RESULT_MISSING_TOTAL_LENGTH_FOR_FIRST_FRAGMENT_1D11374F = + "Missing totalLength for the first fragment of a chunked TsBlock"; + public static final String + EXCEPTION_CHUNKED_RESULT_ASSEMBLED_LENGTH_ARG_DOES_NOT_MATCH_TOTAL_LENGTH_ARG_925DCB15 = + "Assembled fragment length %d does not match the total length %d"; } diff --git a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java index fc059e26f5f32..0f448debed5ec 100644 --- a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -108,4 +108,15 @@ private RpcMessages() {} "通过列名 %s 获取的值为 NULL。"; public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 = "不支持数据类型 %s。"; + + public static final String EXCEPTION_CHUNKED_RESULT_EMPTY_FRAGMENT_03A0EF39 = + "拼装分片查询结果时收到空分片"; + public static final String + EXCEPTION_CHUNKED_RESULT_FRAGMENT_OFFSET_ARG_NOT_GREATER_THAN_ACCUMULATED_LENGTH_ARG_A0F032E3 = + "分片 offset %d 不大于已累计的长度 %d"; + public static final String EXCEPTION_CHUNKED_RESULT_MISSING_TOTAL_LENGTH_FOR_FIRST_FRAGMENT_1D11374F = + "分片 TsBlock 的首个分片缺少 totalLength"; + public static final String + EXCEPTION_CHUNKED_RESULT_ASSEMBLED_LENGTH_ARG_DOES_NOT_MATCH_TOTAL_LENGTH_ARG_925DCB15 = + "拼装后的分片长度 %d 与 totalLength %d 不一致"; } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/IoTDBRpcDataSet.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/IoTDBRpcDataSet.java index 32f5ea5d45a32..a88f473649f9f 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/IoTDBRpcDataSet.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/IoTDBRpcDataSet.java @@ -34,6 +34,7 @@ import org.apache.tsfile.read.common.type.Type; import org.apache.tsfile.utils.Binary; import org.apache.tsfile.utils.DateUtils; +import org.apache.tsfile.utils.PublicBAOS; import java.nio.ByteBuffer; import java.sql.Timestamp; @@ -84,6 +85,8 @@ public class IoTDBRpcDataSet { private int queryResultIndex; // the index of bytebuffer in queryResult private int tsBlockSize; // the size of current tsBlock private int tsBlockIndex; // the row index in current tsBlock + // the fragment of a TsBlock that is partially received; null if there is no pending fragment + private ResultFragmentBuffer fragmentBuffer; private final ZoneId zoneId; private final String timeFormat; @@ -111,6 +114,53 @@ public IoTDBRpcDataSet( int timeFactor, boolean tableModel, List columnIndex2TsBlockColumnIndexList) { + this( + sql, + columnNameList, + columnTypeList, + columnNameIndex, + ignoreTimeStamp, + moreData, + queryId, + statementId, + client, + sessionId, + queryResult, + fetchSize, + timeout, + zoneId, + timeFormat, + timeFactor, + tableModel, + columnIndex2TsBlockColumnIndexList, + false, + 0, + 0); + } + + @SuppressWarnings({"squid:S3776", "squid:S107"}) // Suppress high Cognitive Complexity warning + public IoTDBRpcDataSet( + String sql, + List columnNameList, + List columnTypeList, + Map columnNameIndex, + boolean ignoreTimeStamp, + boolean moreData, + long queryId, + long statementId, + IClientRPCService.Iface client, + long sessionId, + List queryResult, + int fetchSize, + long timeout, + ZoneId zoneId, + String timeFormat, + int timeFactor, + boolean tableModel, + List columnIndex2TsBlockColumnIndexList, + boolean offsetSet, + int offset, + int totalLength) { this.sessionId = sessionId; this.statementId = statementId; // only used for tree model, table model this field will always be true @@ -184,12 +234,13 @@ public IoTDBRpcDataSet( } } - this.queryResult = queryResult; + this.queryResult = new ArrayList<>(); + this.queryResultIndex = 0; this.queryResultSize = 0; if (queryResult != null) { - queryResultSize = queryResult.size(); + addResponse(queryResult, offsetSet, offset, totalLength); + this.queryResultSize = this.queryResult.size(); } - this.queryResultIndex = 0; this.tsBlockSize = 0; this.tsBlockIndex = -1; this.zoneId = zoneId; @@ -227,6 +278,7 @@ public void close() throws StatementExecutionException, TException { } client = null; isClosed = true; + fragmentBuffer = null; } public boolean next() throws StatementExecutionException, IoTDBConnectionException { @@ -241,7 +293,13 @@ public boolean next() throws StatementExecutionException, IoTDBConnectionExcepti return true; } - if (moreData && fetchResults() && hasCachedByteBuffer()) { + // one oversized TsBlock may need several fetch RPCs before it is fully assembled + while (moreData && !hasCachedByteBuffer()) { + if (!fetchResults()) { + break; + } + } + if (hasCachedByteBuffer()) { constructOneTsBlock(); constructOneRow(); return true; @@ -262,6 +320,9 @@ public boolean fetchResults() throws StatementExecutionException, IoTDBConnectio TSFetchResultsReq req = new TSFetchResultsReq(sessionId, sql, fetchSize, queryId, true); req.setStatementId(statementId); req.setTimeout(timeout); + // a new client always sets offset (0 if nothing pending), which enables bounded-payload + // responses on the server + req.setOffset(fragmentBuffer == null ? 0 : fragmentBuffer.getNextOffset()); try { TSFetchResultsResp resp = client.fetchResultsV2(req); RpcUtils.verifySuccess(resp.getStatus()); @@ -269,12 +330,10 @@ public boolean fetchResults() throws StatementExecutionException, IoTDBConnectio if (!resp.hasResultSet) { close(); } else { - queryResult = resp.getQueryResult(); + queryResult = new ArrayList<>(); queryResultIndex = 0; - queryResultSize = 0; - if (queryResult != null) { - queryResultSize = queryResult.size(); - } + addResponse(resp.getQueryResult(), resp.isSetOffset(), resp.offset, resp.totalLength); + queryResultSize = queryResult.size(); this.tsBlockSize = 0; this.tsBlockIndex = -1; } @@ -286,6 +345,98 @@ public boolean fetchResults() throws StatementExecutionException, IoTDBConnectio } } + /** + * Adds the blocks of one response to {@link #queryResult}. If the response ends with a fragment, + * the fragment is either appended to the pending {@link ResultFragmentBuffer} or starts a new + * one; the corresponding TsBlock is only added to {@link #queryResult} once all of its bytes have + * been received. + */ + private void addResponse( + List blocks, boolean offsetSet, int nextOffset, int totalLength) { + if (blocks == null || blocks.isEmpty()) { + return; + } + int blockIndex = 0; + + if (fragmentBuffer != null) { + fragmentBuffer.append(blocks.get(blockIndex++)); + if (offsetSet && blockIndex == blocks.size()) { + fragmentBuffer.updateOffset(nextOffset); + return; + } + queryResult.add(fragmentBuffer.toByteBuffer()); + fragmentBuffer = null; + } + + int lastCompleteBlockIndex = offsetSet ? blocks.size() - 1 : blocks.size(); + while (blockIndex < lastCompleteBlockIndex) { + queryResult.add(blocks.get(blockIndex++)); + } + + if (offsetSet) { + if (totalLength <= 0) { + throw new IllegalArgumentException( + RpcMessages.EXCEPTION_CHUNKED_RESULT_MISSING_TOTAL_LENGTH_FOR_FIRST_FRAGMENT_1D11374F); + } + fragmentBuffer = new ResultFragmentBuffer(totalLength); + fragmentBuffer.append(blocks.get(blockIndex)); + fragmentBuffer.updateOffset(nextOffset); + } + } + + /** Accumulates the fragments of one oversized serialized TsBlock until it is complete. */ + private static class ResultFragmentBuffer { + + private final PublicBAOS buffer; + private final int totalLength; + private int nextOffset; + + private ResultFragmentBuffer(int totalLength) { + this.buffer = new PublicBAOS(totalLength); + this.totalLength = totalLength; + } + + private int getNextOffset() { + return nextOffset; + } + + private void append(ByteBuffer fragment) { + if (!fragment.hasRemaining()) { + throw new IllegalArgumentException( + RpcMessages.EXCEPTION_CHUNKED_RESULT_EMPTY_FRAGMENT_03A0EF39); + } + // a fragment may be a slice of a larger buffer, so honor arrayOffset and position + buffer.write( + fragment.array(), fragment.arrayOffset() + fragment.position(), fragment.remaining()); + } + + private void updateOffset(int nextOffset) { + // the accumulated buffer only grows with non-empty fragments, so this also rejects a + // non-increasing offset + if (nextOffset != buffer.size()) { + throw new IllegalArgumentException( + String.format( + RpcMessages + .EXCEPTION_CHUNKED_RESULT_FRAGMENT_OFFSET_ARG_NOT_GREATER_THAN_ACCUMULATED_LENGTH_ARG_A0F032E3, + nextOffset, + buffer.size())); + } + this.nextOffset = nextOffset; + } + + private ByteBuffer toByteBuffer() { + if (buffer.size() != totalLength) { + throw new IllegalArgumentException( + String.format( + RpcMessages + .EXCEPTION_CHUNKED_RESULT_ASSEMBLED_LENGTH_ARG_DOES_NOT_MATCH_TOTAL_LENGTH_ARG_925DCB15, + buffer.size(), + totalLength)); + } + return ByteBuffer.wrap(buffer.getBuf(), 0, buffer.size()); + } + } + public boolean hasCachedBlock() { return (curTsBlock != null && tsBlockIndex < tsBlockSize - 1); } diff --git a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/IoTDBRpcDataSetFragmentTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/IoTDBRpcDataSetFragmentTest.java new file mode 100644 index 0000000000000..c68553cb013f9 --- /dev/null +++ b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/IoTDBRpcDataSetFragmentTest.java @@ -0,0 +1,249 @@ +/* + * 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.rpc; + +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSCloseOperationReq; +import org.apache.iotdb.service.rpc.thrift.TSFetchResultsReq; +import org.apache.iotdb.service.rpc.thrift.TSFetchResultsResp; + +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.TsBlockSerde; +import org.junit.Test; +import org.mockito.ArgumentCaptor; + +import java.nio.ByteBuffer; +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class IoTDBRpcDataSetFragmentTest { + + private static final String COLUMN = "s1"; + + @Test + public void testAssembleOneOversizedTsBlockAcrossResponses() throws Exception { + byte[] bytes = serializeInt64Block(); + int totalLength = bytes.length; + int split = totalLength / 2; + + IClientRPCService.Iface client = mockClient(); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenReturn( + fragmentResponse(ByteBuffer.wrap(bytes, 0, split), split, totalLength, true), + fragmentResponse(ByteBuffer.wrap(bytes, split, totalLength - split), -1, -1, false)); + + IoTDBRpcDataSet dataSet = newDataSet(client); + + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertFalse(dataSet.next()); + + ArgumentCaptor captor = ArgumentCaptor.forClass(TSFetchResultsReq.class); + verify(client, times(2)).fetchResultsV2(captor.capture()); + assertEquals(0, captor.getAllValues().get(0).getOffset()); + assertEquals(split, captor.getAllValues().get(1).getOffset()); + } + + @Test + public void testCompleteBlocksAndTrailingFragmentInOneResponse() throws Exception { + byte[] first = serializeInt64Block(); + byte[] second = serializeInt64Block(); + int split = second.length / 2; + + IClientRPCService.Iface client = mockClient(); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenReturn( + response( + List.of(ByteBuffer.wrap(first), ByteBuffer.wrap(second, 0, split)), + split, + second.length, + true), + fragmentResponse(ByteBuffer.wrap(second, split, second.length - split), -1, -1, false)); + + IoTDBRpcDataSet dataSet = newDataSet(client); + + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertFalse(dataSet.next()); + } + + @Test + public void testLegacyResponseWithoutOffset() throws Exception { + byte[] bytes = serializeInt64Block(); + IClientRPCService.Iface client = mockClient(); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenReturn(response(List.of(ByteBuffer.wrap(bytes)), -1, -1, false)); + + IoTDBRpcDataSet dataSet = newDataSet(client); + + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertFalse(dataSet.next()); + } + + @Test + public void testBlocksInExecuteResponseAreConsumedBeforeFetching() throws Exception { + byte[] bytes = serializeInt64Block(); + IClientRPCService.Iface client = mockClient(); + IoTDBRpcDataSet dataSet = + new IoTDBRpcDataSet( + "", + List.of(COLUMN), + List.of(TSDataType.INT64.name()), + null, + true, + false, + 1L, + 2L, + client, + 3L, + List.of(ByteBuffer.wrap(bytes)), + 100, + 0, + ZoneId.systemDefault(), + RpcUtils.DEFAULT_TIME_FORMAT, + 1_000_000, + true, + new ArrayList<>(List.of(0))); + + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertFalse(dataSet.next()); + verify(client, times(0)).fetchResultsV2(any(TSFetchResultsReq.class)); + } + + @Test(expected = IllegalArgumentException.class) + public void testRejectMissingTotalLengthForFirstFragment() throws Exception { + byte[] bytes = serializeInt64Block(); + IClientRPCService.Iface client = mockClient(); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenReturn( + fragmentResponse( + ByteBuffer.wrap(bytes, 0, bytes.length / 2), bytes.length / 2, -1, true)); + + IoTDBRpcDataSet dataSet = newDataSet(client); + dataSet.next(); + } + + @Test(expected = IllegalArgumentException.class) + public void testRejectNonIncreasingFragmentOffset() throws Exception { + byte[] bytes = serializeInt64Block(); + int split = bytes.length / 2; + IClientRPCService.Iface client = mockClient(); + // the offset must equal the accumulated length, otherwise the dataset must refuse to continue + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenReturn( + fragmentResponse(ByteBuffer.wrap(bytes, 0, split), split - 1, bytes.length, true)); + + IoTDBRpcDataSet dataSet = newDataSet(client); + dataSet.next(); + } + + private static IoTDBRpcDataSet newDataSet(IClientRPCService.Iface client) { + return new IoTDBRpcDataSet( + "", + List.of(COLUMN), + List.of(TSDataType.INT64.name()), + null, + true, + true, + 1L, + 2L, + client, + 3L, + null, + 100, + 0, + ZoneId.systemDefault(), + RpcUtils.DEFAULT_TIME_FORMAT, + 1_000_000, + true, + new ArrayList<>(List.of(0))); + } + + private static IClientRPCService.Iface mockClient() throws Exception { + IClientRPCService.Iface client = mock(IClientRPCService.Iface.class); + when(client.closeOperation(any(TSCloseOperationReq.class))) + .thenReturn(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + return client; + } + + private static TSFetchResultsResp fragmentResponse( + ByteBuffer fragment, int nextOffset, int totalLength, boolean moreData) { + return response(Collections.singletonList(fragment), nextOffset, totalLength, moreData); + } + + private static TSFetchResultsResp response( + List blocks, int nextOffset, int totalLength, boolean moreData) { + TSFetchResultsResp resp = + new TSFetchResultsResp(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS), true, true); + resp.setQueryResult(blocks); + resp.setMoreData(moreData); + if (nextOffset >= 0) { + resp.setOffset(nextOffset); + } + if (totalLength >= 0) { + resp.setTotalLength(totalLength); + } + return resp; + } + + /** [time=1, s1=10], [time=2, s1=20] as an INT64 column. */ + private static byte[] serializeInt64Block() throws Exception { + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT64)); + builder.getTimeColumnBuilder().writeLong(1L); + builder.getColumnBuilder(0).writeLong(10L); + builder.declarePosition(); + builder.getTimeColumnBuilder().writeLong(2L); + builder.getColumnBuilder(0).writeLong(20L); + builder.declarePosition(); + TsBlock tsBlock = builder.build(); + + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + byte[] bytes = new byte[serialized.remaining()]; + serialized.duplicate().get(bytes); + return bytes; + } +} diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java index ca8862cf1c989..089216aa6e7ab 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java @@ -104,6 +104,7 @@ public class SessionConnection { private static final Logger logger = LoggerFactory.getLogger(SessionConnection.class); private static final String USE_ENCRYPTED_PASSWORD_KEY = "use_encrypted_password"; private static final String ENABLE_AUTO_CREATE_SCHEMA_KEY = "enable_auto_create_schema"; + private static final String ENABLE_RESULT_CHUNKING_KEY = "enable_result_chunking"; public static final String MSG_RECONNECTION_FAIL = SessionMessages.EXCEPTION_FAIL_TO_RECONNECT_TO_SERVER_PLEASE_CHECK_SERVER_STATUS_D9A2790F; protected Session session; @@ -252,6 +253,8 @@ private void init( openReq.setZoneId(zoneId.toString()); openReq.putToConfiguration("version", session.version.toString()); openReq.putToConfiguration("sql_dialect", sqlDialect); + // this client can assemble fragmented query results, see TSFetchResultsReq#offset + openReq.putToConfiguration(ENABLE_RESULT_CHUNKING_KEY, Boolean.TRUE.toString()); if (session.useEncryptedPassword) { openReq.putToConfiguration(USE_ENCRYPTED_PASSWORD_KEY, Boolean.TRUE.toString()); } @@ -498,7 +501,10 @@ protected SessionDataSet executeQueryStatement(String sql, long timeout) zoneId, timeFactor, execResp.isSetTableModel() && execResp.isTableModel(), - execResp.getColumnIndex2TsBlockColumnIndexList()); + execResp.getColumnIndex2TsBlockColumnIndexList(), + execResp.isSetOffset(), + execResp.offset, + execResp.totalLength); } protected void executeNonQueryStatement(String sql) @@ -567,7 +573,10 @@ protected SessionDataSet executeRawDataQuery( zoneId, timeFactor, execResp.isSetTableModel() && execResp.isTableModel(), - execResp.getColumnIndex2TsBlockColumnIndexList()); + execResp.getColumnIndex2TsBlockColumnIndexList(), + execResp.isSetOffset(), + execResp.offset, + execResp.totalLength); } protected SessionDataSet executeLastDataQueryForOnePrefixPath(final List prefixes) @@ -607,7 +616,10 @@ protected SessionDataSet executeLastDataQueryForOnePrefixPath(final List zoneId, timeFactor, false, - tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList()); + tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList(), + tsExecuteStatementResp.isSetOffset(), + tsExecuteStatementResp.offset, + tsExecuteStatementResp.totalLength); } protected Pair executeLastDataQueryForOneDevice( @@ -657,7 +669,10 @@ protected Pair executeLastDataQueryForOneDevice( zoneId, timeFactor, tsExecuteStatementResp.isSetTableModel() && tsExecuteStatementResp.isTableModel(), - tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList()), + tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList(), + tsExecuteStatementResp.isSetOffset(), + tsExecuteStatementResp.offset, + tsExecuteStatementResp.totalLength), redirectedEndPoint); } @@ -700,7 +715,10 @@ protected SessionDataSet executeLastDataQuery(List paths, long time, lon zoneId, timeFactor, tsExecuteStatementResp.isSetTableModel() && tsExecuteStatementResp.isTableModel(), - tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList()); + tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList(), + tsExecuteStatementResp.isSetOffset(), + tsExecuteStatementResp.offset, + tsExecuteStatementResp.totalLength); } protected SessionDataSet executeAggregationQuery( @@ -782,7 +800,10 @@ private SessionDataSet executeAggregationQuery(TSAggregationQueryReq tsAggregati zoneId, timeFactor, tsExecuteStatementResp.isSetTableModel() && tsExecuteStatementResp.isTableModel(), - tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList()); + tsExecuteStatementResp.getColumnIndex2TsBlockColumnIndexList(), + tsExecuteStatementResp.isSetOffset(), + tsExecuteStatementResp.offset, + tsExecuteStatementResp.totalLength); } private TSAggregationQueryReq createAggregationQueryReq( diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java index e15320c1a1519..dd3a374c21b93 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java @@ -147,6 +147,8 @@ public void init( openReq.setPassword(password); openReq.setZoneId(zoneId.toString()); openReq.putToConfiguration("version", version); + // this client can assemble fragmented query results, see TSFetchResultsReq#offset + openReq.putToConfiguration("enable_result_chunking", Boolean.TRUE.toString()); try { TSOpenSessionResp openResp = client.openSession(openReq); @@ -210,7 +212,10 @@ protected SessionDataSet executeQueryStatement(String sql, long timeout, int fet zoneId, timeFactor, execResp.isSetTableModel() && execResp.isTableModel(), - execResp.getColumnIndex2TsBlockColumnIndexList()); + execResp.getColumnIndex2TsBlockColumnIndexList(), + execResp.isSetOffset(), + execResp.offset, + execResp.totalLength); } /** Create outside the publication lock: TLS setup itself can perform network I/O. */ diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 574feb482cc56..3a38ae3c1cc6d 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -4971,4 +4971,11 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ARG_ARG_MS_ARG_INVOCATIONS_ARG_APPLICATIONS_F056BD84 = "%s: %s ms, %s invocations, %s applications"; + + public static final String + EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_DOES_NOT_MATCH_PENDING_FRAGMENT_OFFSET_ARG_A1B13DBC = + "Request offset %d does not match the pending result fragment offset %d"; + public static final String + EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_IS_INVALID_WITHOUT_PENDING_FRAGMENT_C403454F = + "Request offset %d is invalid because no result fragment is pending"; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 3895194221b60..d47d0514a8f15 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -5671,4 +5671,11 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ARG_ARG_MS_ARG_INVOCATIONS_ARG_APPLICATIONS_F056BD84 = "%s：%s ms，调用 %s 次，应用 %s 次"; + + public static final String + EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_DOES_NOT_MATCH_PENDING_FRAGMENT_OFFSET_ARG_A1B13DBC = + "请求 offset %d 与待发送结果分片的 offset %d 不一致"; + public static final String + EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_IS_INVALID_WITHOUT_PENDING_FRAGMENT_C403454F = + "请求 offset %d 无效：当前没有待发送的结果分片"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java index 590e87044cb9d..26e0252a19b35 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java @@ -66,6 +66,9 @@ public abstract class IClientSession { private boolean enableAutoCreateSchema = true; + // set by clients that can assemble fragmented query results, see TSFetchResultsReq#offset + private boolean enableResultChunking = false; + private long lastActiveTime = CommonDateTimeUtils.currentTime(); public abstract String getClientAddress(); @@ -224,6 +227,14 @@ public void setEnableAutoCreateSchema(boolean enableAutoCreateSchema) { this.enableAutoCreateSchema = enableAutoCreateSchema; } + public boolean isEnableResultChunking() { + return enableResultChunking; + } + + public void setEnableResultChunking(boolean enableResultChunking) { + this.enableResultChunking = enableResultChunking; + } + /** * Add a prepared statement to this session. * diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index e6eedc4b636a9..67d59b59ddfe7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -274,6 +274,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler { public static final String ERROR_CODE = DataNodeMiscMessages.ERROR_CODE; private static final String USE_ENCRYPTED_PASSWORD_KEY = "use_encrypted_password"; private static final String ENABLE_AUTO_CREATE_SCHEMA_KEY = "enable_auto_create_schema"; + private static final String ENABLE_RESULT_CHUNKING_KEY = "enable_result_chunking"; private static final TSProtocolVersion CURRENT_RPC_VERSION = TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3; @@ -303,25 +304,76 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler { @FunctionalInterface public interface SelectResult { - boolean apply(TSExecuteStatementResp resp, IQueryExecution queryExecution, int fetchSize) + boolean apply( + TSExecuteStatementResp resp, + IQueryExecution queryExecution, + int fetchSize, + int maxPayloadSize) throws IoTDBException, IOException; } private static final SelectResult SELECT_RESULT = - (resp, queryExecution, fetchSize) -> { - Pair, Boolean> pair = - QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, fetchSize); - resp.setQueryResult(pair.left); - return pair.right; - }; + (resp, queryExecution, fetchSize, maxPayloadSize) -> + fillQueryResult(resp, fetchResultChunk(queryExecution, fetchSize, 0, maxPayloadSize)); private static final SelectResult OLD_SELECT_RESULT = - (resp, queryExecution, fetchSize) -> { + (resp, queryExecution, fetchSize, maxPayloadSize) -> { Pair pair = convertTsBlockByFetchSize(queryExecution, fetchSize); resp.setQueryDataSet(pair.left); return pair.right; }; + /** + * Fetches one response-worth of query results. Fragmentation is disabled when {@code + * maxPayloadSize} is not positive, i.e. the client cannot assemble fragments. + */ + private static QueryDataSetUtils.ResultChunk fetchResultChunk( + IQueryExecution queryExecution, int fetchSize, int requestOffset, int maxPayloadSize) + throws IoTDBException { + if (maxPayloadSize > 0) { + return QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, fetchSize, requestOffset, maxPayloadSize); + } + Pair, Boolean> pair = + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, fetchSize); + return QueryDataSetUtils.ResultChunk.complete(pair.left, pair.right); + } + + private static boolean fillQueryResult( + TSExecuteStatementResp resp, QueryDataSetUtils.ResultChunk chunk) { + resp.setQueryResult(chunk.getQueryResult()); + if (chunk.hasNextOffset()) { + resp.setOffset(chunk.getNextOffset()); + if (chunk.hasTotalLength()) { + resp.setTotalLength(chunk.getTotalLength()); + } + } + return chunk.isFinished(); + } + + private static boolean fillQueryResult( + TSFetchResultsResp resp, QueryDataSetUtils.ResultChunk chunk) { + resp.setQueryResult(chunk.getQueryResult()); + if (chunk.hasNextOffset()) { + resp.setOffset(chunk.getNextOffset()); + if (chunk.hasTotalLength()) { + resp.setTotalLength(chunk.getTotalLength()); + } + } + return chunk.isFinished(); + } + + /** + * Returns the maximum payload size of one response carrying query results, or 0 if the session's + * client cannot assemble fragmented results. The payload bound is shared with the MPP data + * exchange RPC response, see mpp_data_exchange_max_payload_size_in_bytes. + */ + private static int getMaxResultPayloadSize(IClientSession clientSession) { + return clientSession.isEnableResultChunking() + ? config.getMppDataExchangeMaxPayloadSizeInBytes() + : 0; + } + public ClientRPCServiceImpl() { partitionFetcher = ClusterPartitionFetcher.getInstance(); schemaFetcher = ClusterSchemaFetcher.getInstance(); @@ -495,7 +547,12 @@ private TSExecuteStatementResp executeStatementInternal( statementType = statementType == null ? StatementType.QUERY : statementType; resp = createResponse(queryExecution.getDatasetHeader(), queryId); resp.setStatus(result.status); - finished = setResult.apply(resp, queryExecution, request.getFetchSize()); + finished = + setResult.apply( + resp, + queryExecution, + request.getFetchSize(), + getMaxResultPayloadSize(clientSession)); resp.setMoreData(!finished); if (quota != null) { quota.addReadResult(resp.getQueryResult()); @@ -768,7 +825,9 @@ private TSExecuteStatementResp executeRawDataQueryInternal( if (queryExecution.isQuery()) { resp = createResponse(queryExecution.getDatasetHeader(), queryId); resp.setStatus(result.status); - finished = setResult.apply(resp, queryExecution, req.fetchSize); + finished = + setResult.apply( + resp, queryExecution, req.fetchSize, getMaxResultPayloadSize(clientSession)); resp.setMoreData(!finished); quota.addReadResult(resp.getQueryResult()); } else { @@ -864,7 +923,9 @@ private TSExecuteStatementResp executeLastDataQueryInternal( if (queryExecution.isQuery()) { resp = createResponse(queryExecution.getDatasetHeader(), queryId); resp.setStatus(result.status); - finished = setResult.apply(resp, queryExecution, req.fetchSize); + finished = + setResult.apply( + resp, queryExecution, req.fetchSize, getMaxResultPayloadSize(clientSession)); resp.setMoreData(!finished); quota.addReadResult(resp.getQueryResult()); } else { @@ -962,7 +1023,9 @@ private TSExecuteStatementResp executeAggregationQueryInternal( if (queryExecution.isQuery()) { resp = createResponse(queryExecution.getDatasetHeader(), queryId); resp.setStatus(result.status); - finished = setResult.apply(resp, queryExecution, req.fetchSize); + finished = + setResult.apply( + resp, queryExecution, req.fetchSize, getMaxResultPayloadSize(clientSession)); resp.setMoreData(!finished); quota.addReadResult(resp.getQueryResult()); } else { @@ -1455,7 +1518,9 @@ public TSExecuteStatementResp executeFastLastDataQueryForOneDeviceV2( regionReplicaSets.get(regionReplicaSets.size() - 1).dataNodeLocations.get(0), clientSession.isInternalRpcSession())); resp.setStatus(tsstatus); - finished = SELECT_RESULT.apply(resp, queryExecution, req.fetchSize); + finished = + SELECT_RESULT.apply( + resp, queryExecution, req.fetchSize, getMaxResultPayloadSize(clientSession)); resp.setMoreData(!finished); quota.addReadResult(resp.getQueryResult()); } else { @@ -1570,14 +1635,15 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) { statementType = queryExecution.getStatementType(); try (SetThreadName queryName = new SetThreadName(queryExecution.getQueryId())) { - Pair, Boolean> pair = - QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, req.fetchSize); - List result = pair.left; - finished = pair.right; - boolean hasResultSet = !result.isEmpty(); - resp.setHasResultSet(hasResultSet); + QueryDataSetUtils.ResultChunk chunk = + fetchResultChunk( + queryExecution, + req.fetchSize, + req.isSetOffset() ? req.getOffset() : 0, + req.isSetOffset() ? config.getMppDataExchangeMaxPayloadSizeInBytes() : 0); + finished = fillQueryResult(resp, chunk); + resp.setHasResultSet(!chunk.getQueryResult().isEmpty()); resp.setIsAlign(true); - resp.setQueryResult(result); resp.setMoreData(!finished); return resp; } @@ -1634,6 +1700,7 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { clientSession.setDatabaseName(database.get().toLowerCase(Locale.ENGLISH)); } clientSession.setEnableAutoCreateSchema(parseEnableAutoCreateSchema(req)); + clientSession.setEnableResultChunking(parseEnableResultChunking(req)); } TSOpenSessionResp resp = new TSOpenSessionResp(tsStatus, CURRENT_RPC_VERSION); Map configuration = new HashMap<>(); @@ -1708,6 +1775,12 @@ private boolean parseEnableAutoCreateSchema(TSOpenSessionReq req) { || Boolean.parseBoolean(configuration.get(ENABLE_AUTO_CREATE_SCHEMA_KEY)); } + private boolean parseEnableResultChunking(TSOpenSessionReq req) { + Map configuration = req.configuration; + return configuration != null + && Boolean.parseBoolean(configuration.get(ENABLE_RESULT_CHUNKING_KEY)); + } + @Override public TSStatus closeSession(TSCloseSessionReq req) { return new TSStatus( @@ -1797,7 +1870,8 @@ public TSPrepareResp prepareStatement(TSPrepareReq req) { @Override public TSExecuteStatementResp executePreparedStatement(TSExecutePreparedReq req) { - return executeStatementInternal(new TSExecutePreparedReqAdapter(req), setResultForPrepared); + // Use V2 format (queryResult) to match IoTDBTablePreparedStatement client + return executeStatementInternal(new TSExecutePreparedReqAdapter(req), SELECT_RESULT); } @Override @@ -1818,15 +1892,6 @@ public TSStatus deallocatePreparedStatement(TSDeallocatePreparedReq req) { } } - private final SelectResult setResultForPrepared = - (resp, queryExecution, fetchSize) -> { - // Use V2 format (queryResult) to match IoTDBTablePreparedStatement client - Pair, Boolean> pair = - QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, fetchSize); - resp.setQueryResult(pair.left); - return pair.right; - }; - // ========================= End PreparedStatement RPC Methods ========================= @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java index 32e5557b8f922..210c2ea48008d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java @@ -50,9 +50,19 @@ public interface ISourceHandle extends Accountable { * iterator with receive(). When one of these two methods is called, the cursor in iterator will * forward. * + * @param retainMemory if true, the memory that the returned block retains in the memory pool is + * not freed when the block is taken out of this handle. The caller must release it via {@link + * #freeRetainedMemory()} once the block is no longer needed. Implementations whose serialized + * blocks are not reserved from the memory pool ignore this flag. * @throws IoTDBException throw IoTDBException */ - ByteBuffer getSerializedTsBlock() throws IoTDBException; + ByteBuffer getSerializedTsBlock(boolean retainMemory) throws IoTDBException; + + /** + * Releases the memory retained by the blocks taken out with {@code retainMemory=true} and not + * released yet. It is safe to call when nothing is retained, and after abort/close. + */ + void freeRetainedMemory(); /** If there are more TsBlocks. */ boolean isFinished(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java index 544421712669a..94f6e09fa058c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java @@ -153,7 +153,8 @@ public TsBlock receive() { } @Override - public ByteBuffer getSerializedTsBlock() throws IoTDBException { + public ByteBuffer getSerializedTsBlock(boolean retainMemory) throws IoTDBException { + // the returned bytes are a fresh serialization, not reserved from the memory pool TsBlock tsBlock = receive(); if (LOGGER.isDebugEnabled()) { LOGGER.debug(DataNodeQueryMessages.GET_SERIALIZED_TSBLOCK, CommonUtils.toString(tsBlock)); @@ -174,6 +175,11 @@ public ByteBuffer getSerializedTsBlock() throws IoTDBException { } } + @Override + public void freeRetainedMemory() { + // nothing is reserved from the memory pool + } + @Override public boolean isFinished() { synchronized (queue) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 8ee7ad601ef61..703ce3407f6a3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -109,6 +109,13 @@ public class SourceHandle implements ISourceHandle { /** The actual buffered memory in bytes, including the amount of memory being reserved. */ private long bufferRetainedSizeInBytes = 0L; + /** + * The memory in bytes taken out of the buffer with {@code retainMemory=true}, still reserved in + * the memory pool on behalf of the caller. Freed by {@link #freeRetainedMemory()}, or by + * abort/close as a fallback. + */ + private long pendingRetainedSizeInBytes = 0L; + private int currSequenceId = 0; private int nextSequenceId = 0; private int lastSequenceId = Integer.MAX_VALUE; @@ -232,7 +239,7 @@ public SourceHandle( @Override public synchronized TsBlock receive() { - ByteBuffer tsBlock = getSerializedTsBlock(); + ByteBuffer tsBlock = getSerializedTsBlock(false); if (tsBlock != null) { long startTime = System.nanoTime(); try { @@ -247,7 +254,7 @@ public synchronized TsBlock receive() { } @Override - public synchronized ByteBuffer getSerializedTsBlock() { + public synchronized ByteBuffer getSerializedTsBlock(boolean retainMemory) { long startTime = System.nanoTime(); try (SetThreadName sourceHandleName = new SetThreadName(threadName)) { checkState(); @@ -270,13 +277,17 @@ public synchronized ByteBuffer getSerializedTsBlock() { currSequenceId += 1; if (retainedSize > 0) { bufferRetainedSizeInBytes -= retainedSize; - localMemoryManager - .getQueryPool() - .free( - localFragmentInstanceId.getQueryId(), - fullFragmentInstanceId, - localPlanNodeId, - retainedSize); + if (retainMemory) { + pendingRetainedSizeInBytes += retainedSize; + } else { + localMemoryManager + .getQueryPool() + .free( + localFragmentInstanceId.getQueryId(), + fullFragmentInstanceId, + localPlanNodeId, + retainedSize); + } } if (sequenceIdToTsBlock.isEmpty() && !isFinished()) { @@ -296,6 +307,22 @@ public synchronized ByteBuffer getSerializedTsBlock() { } } + @Override + public synchronized void freeRetainedMemory() { + if (pendingRetainedSizeInBytes <= 0) { + return; + } + long retainedSize = pendingRetainedSizeInBytes; + pendingRetainedSizeInBytes = 0; + localMemoryManager + .getQueryPool() + .free( + localFragmentInstanceId.getQueryId(), + fullFragmentInstanceId, + localPlanNodeId, + retainedSize); + } + private synchronized void trySubmitGetDataBlocksTask() { if (aborted || closed) { return; @@ -444,16 +471,7 @@ public synchronized void abort() { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } sequenceIdToDataBlockSize.clear(); - if (bufferRetainedSizeInBytes > 0) { - localMemoryManager - .getQueryPool() - .free( - localFragmentInstanceId.getQueryId(), - fullFragmentInstanceId, - localPlanNodeId, - bufferRetainedSizeInBytes); - bufferRetainedSizeInBytes = 0; - } + freeAllRetainedMemory(); aborted = true; sourceHandleListener.onAborted(this); } @@ -472,16 +490,7 @@ public synchronized void abort(Throwable t) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } sequenceIdToDataBlockSize.clear(); - if (bufferRetainedSizeInBytes > 0) { - localMemoryManager - .getQueryPool() - .free( - localFragmentInstanceId.getQueryId(), - fullFragmentInstanceId, - localPlanNodeId, - bufferRetainedSizeInBytes); - bufferRetainedSizeInBytes = 0; - } + freeAllRetainedMemory(); aborted = true; sourceHandleListener.onAborted(this); } @@ -500,16 +509,7 @@ public synchronized void close() { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } sequenceIdToDataBlockSize.clear(); - if (bufferRetainedSizeInBytes > 0) { - localMemoryManager - .getQueryPool() - .free( - localFragmentInstanceId.getQueryId(), - fullFragmentInstanceId, - localPlanNodeId, - bufferRetainedSizeInBytes); - bufferRetainedSizeInBytes = 0; - } + freeAllRetainedMemory(); closed = true; executorService.submit(new SendCloseSinkChannelEventTask()); currSequenceId = lastSequenceId + 1; @@ -517,6 +517,25 @@ public synchronized void close() { } } + /** + * Frees the buffered memory and the memory handed over with {@code retainMemory=true}, as a + * fallback for blocks the caller has not released yet. + */ + private void freeAllRetainedMemory() { + long totalRetainedSize = bufferRetainedSizeInBytes + pendingRetainedSizeInBytes; + if (totalRetainedSize > 0) { + localMemoryManager + .getQueryPool() + .free( + localFragmentInstanceId.getQueryId(), + fullFragmentInstanceId, + localPlanNodeId, + totalRetainedSize); + bufferRetainedSizeInBytes = 0; + pendingRetainedSizeInBytes = 0; + } + } + @Override public boolean isFinished() { return remoteTsBlockedConsumedUp(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java index f4beb09917c27..4a5fa22e1c462 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java @@ -43,7 +43,15 @@ public interface IQueryExecution { Optional getBatchResult() throws IoTDBException; - Optional getByteBufferBatchResult() throws IoTDBException; + /** + * Gets the next serialized TsBlock of the result. If {@code retainMemory} is true, the memory + * that the returned block retains in the memory pool is not freed when the block is taken out of + * the result handle. The caller then hands the block over to {@link #getPendingResultFragment()}, + * which releases the memory once the block has been fully sent. + */ + Optional getByteBufferBatchResult(boolean retainMemory) throws IoTDBException; + + PendingResultFragment getPendingResultFragment(); boolean hasNextResult(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/PendingResultFragment.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/PendingResultFragment.java new file mode 100644 index 0000000000000..21b988d2a77f5 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/PendingResultFragment.java @@ -0,0 +1,68 @@ +/* + * 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.queryengine.plan.execution; + +import java.nio.ByteBuffer; + +/** + * The tail of a serialized TsBlock that did not fit into the previous RPC response and is waiting + * to be sent in the following response(s). The block is already serialized and removed from the + * result handle, so this fragment only holds a reference to it and never copies the payload. + * + *

The memory the block retains stays reserved in the memory pool until the block is fully sent. + * The owner binds a {@code memoryReleaser} which is invoked by {@link #clear()}, so callers only + * interact with this holder and never release the memory themselves. + */ +public class PendingResultFragment { + + private final Runnable memoryReleaser; + + private ByteBuffer block; + private int nextOffset; + + public PendingResultFragment(Runnable memoryReleaser) { + this.memoryReleaser = memoryReleaser; + } + + /** Records that {@code block} has been sent up to {@code nextOffset} bytes. */ + public void reset(ByteBuffer block, int nextOffset) { + this.block = block; + this.nextOffset = nextOffset; + } + + public ByteBuffer getBlock() { + return block; + } + + public int getNextOffset() { + return nextOffset; + } + + /** + * Discards the pending fragment and releases the memory of the blocks that have been fully sent. + * It is also called right after a complete block has been added to a response, when nothing is + * pending. Calling it repeatedly is a no-op. + */ + public void clear() { + block = null; + nextOffset = 0; + memoryReleaser.run(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 5ab3b122e1c6c..d07319e44c641 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -113,6 +113,10 @@ public class QueryExecution implements IQueryExecution { // used for cleaning resultHandle up exactly once private final AtomicBoolean resultHandleCleanUp; + // the serialized TsBlock that is partially sent to the client and waits for the next RPC + private final PendingResultFragment pendingResultFragment = + new PendingResultFragment(this::releaseResultHandleRetainedMemory); + private final UDFClassLoaderManager.QueryScope udfQueryScope; // cost time in ns of finished rpc @@ -549,6 +553,9 @@ private void cleanupForRetry(Throwable t) { */ private void releaseResource(Throwable t, boolean isRetry) { context.recordDeviceEntryDiskIOMetricsOnRelease(); + // the query will not serve any more fetch requests, release the pending fragment's memory and + // make sure a retried execution does not continue a stale fragment + pendingResultFragment.clear(); // close ResultHandle to unblock client's getResult request // Actually, we should not close the ResultHandle when the QueryExecution is Finished. // There are only two scenarios where the ResultHandle should be closed: @@ -717,12 +724,19 @@ private TsBlock getDeserializedTsBlock() { } @Override - public Optional getByteBufferBatchResult() throws IoTDBException { - return getResult(this::getSerializedTsBlock); + public Optional getByteBufferBatchResult(boolean retainMemory) throws IoTDBException { + return getResult(() -> resultHandle.getSerializedTsBlock(retainMemory)); } - private ByteBuffer getSerializedTsBlock() throws IoTDBException { - return resultHandle.getSerializedTsBlock(); + @Override + public PendingResultFragment getPendingResultFragment() { + return pendingResultFragment; + } + + private void releaseResultHandleRetainedMemory() { + if (resultHandle != null) { + resultHandle.freeRetainedMemory(); + } } /** diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java index 148dab7de71fe..621fd329d3b20 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; import org.apache.iotdb.db.queryengine.plan.execution.config.executor.ClusterConfigTaskExecutor; import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor; import org.apache.iotdb.db.queryengine.plan.statement.StatementType; @@ -115,6 +116,9 @@ public class ConfigExecution implements IQueryExecution { private final IConfigTask task; private final IConfigTaskExecutor configTaskExecutor; + // the result set is not reserved from the memory pool, so nothing is released on clear + private final PendingResultFragment pendingResultFragment = new PendingResultFragment(() -> {}); + private final StatementType statementType; private long totalExecutionTime; @@ -277,7 +281,8 @@ public Optional getBatchResult() { } @Override - public Optional getByteBufferBatchResult() throws IoTDBException { + public Optional getByteBufferBatchResult(boolean retainMemory) throws IoTDBException { + // the result set is not reserved from the memory pool, so there is no memory to retain if (!resultSetConsumed) { resultSetConsumed = true; try { @@ -289,6 +294,11 @@ public Optional getByteBufferBatchResult() throws IoTDBException { return Optional.empty(); } + @Override + public PendingResultFragment getPendingResultFragment() { + return pendingResultFragment; + } + // According to the execution process of ConfigExecution, there is only one TsBlock for // this execution. Thus, the hasNextResult will be false once the TsBlock is consumed @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/memory/MemorySourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/memory/MemorySourceHandle.java index 7a4919446cf92..a325fc671a15e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/memory/MemorySourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/memory/MemorySourceHandle.java @@ -77,7 +77,8 @@ public synchronized TsBlock receive() { } @Override - public synchronized ByteBuffer getSerializedTsBlock() throws IoTDBException { + public synchronized ByteBuffer getSerializedTsBlock(boolean retainMemory) throws IoTDBException { + // serialized blocks are not reserved from the memory pool, so retainMemory is ignored hasNext = false; if (result.isEmpty()) { return null; @@ -90,6 +91,11 @@ public synchronized ByteBuffer getSerializedTsBlock() throws IoTDBException { } } + @Override + public void freeRetainedMemory() { + // nothing is reserved from the memory pool + } + @Override public synchronized boolean isFinished() { return !hasNext; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/QueryDataSetUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/QueryDataSetUtils.java index bb40b38dc24b4..a398d2613e922 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/QueryDataSetUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/QueryDataSetUtils.java @@ -20,7 +20,9 @@ package org.apache.iotdb.db.utils; import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; import org.apache.iotdb.service.rpc.thrift.TSQueryDataSet; import org.apache.tsfile.block.column.Column; @@ -251,18 +253,12 @@ public static Pair, Boolean> convertQueryResultByFetchSize( long memorySize = 0; List res = new ArrayList<>(); while (rowCount < fetchSize && memorySize < MAX_RETURN_SIZE) { - Optional optionalByteBuffer = queryExecution.getByteBufferBatchResult(); + Optional optionalByteBuffer = queryExecution.getByteBufferBatchResult(false); if (!optionalByteBuffer.isPresent()) { break; } ByteBuffer byteBuffer = optionalByteBuffer.get(); - byteBuffer.mark(); - int valueColumnCount = byteBuffer.getInt(); - for (int i = 0; i < valueColumnCount; i++) { - byteBuffer.get(); - } - int positionCount = byteBuffer.getInt(); - byteBuffer.reset(); + int positionCount = readPositionCount(byteBuffer); if (positionCount != 0) { res.add(byteBuffer); } @@ -272,6 +268,150 @@ public static Pair, Boolean> convertQueryResultByFetchSize( return new Pair<>(res, !queryExecution.hasNextResult()); } + /** + * Fetches query results with a bounded payload size. A serialized TsBlock larger than {@code + * maxPayloadSize} is cut into fragments: the response only carries the bytes that fit into the + * payload, and the remaining bytes are kept in {@link PendingResultFragment} on the query + * execution, to be sent by the following request carrying the returned offset. + */ + public static ResultChunk convertQueryResultByFetchSize( + IQueryExecution queryExecution, int fetchSize, int requestOffset, int maxPayloadSize) + throws IoTDBException { + fetchSize = fetchSize > 0 ? fetchSize : TSFLE_CONFIG.getMaxTsBlockLineNumber(); + List res = new ArrayList<>(); + int remainingPayloadSize = maxPayloadSize; + int rowCount = 0; + long memorySize = 0; + + PendingResultFragment pendingFragment = queryExecution.getPendingResultFragment(); + ByteBuffer pendingBlock = pendingFragment.getBlock(); + if (pendingBlock != null) { + int nextOffset = pendingFragment.getNextOffset(); + if (requestOffset != nextOffset) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_DOES_NOT_MATCH_PENDING_FRAGMENT_OFFSET_ARG_A1B13DBC, + requestOffset, + nextOffset)); + } + int fragmentSize = Math.min(remainingPayloadSize, pendingBlock.limit() - nextOffset); + res.add(slice(pendingBlock, nextOffset, fragmentSize)); + if (nextOffset + fragmentSize == pendingBlock.limit()) { + pendingFragment.clear(); + // Do not fetch new blocks in the same response, otherwise the response would carry the + // tail of the big block plus up to fetchSize new rows, whose row count is not accounted + // for. One response either continues a fragmented block or delivers new blocks. + return ResultChunk.complete(res, !queryExecution.hasNextResult()); + } + pendingFragment.reset(pendingBlock, nextOffset + fragmentSize); + return ResultChunk.fragment(res, nextOffset + fragmentSize, -1, false); + } else if (requestOffset != 0) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_RESULT_FRAGMENT_REQUEST_OFFSET_ARG_IS_INVALID_WITHOUT_PENDING_FRAGMENT_C403454F, + requestOffset)); + } + + while (rowCount < fetchSize && memorySize < MAX_RETURN_SIZE && remainingPayloadSize > 0) { + Optional optionalByteBuffer = queryExecution.getByteBufferBatchResult(true); + if (!optionalByteBuffer.isPresent()) { + break; + } + ByteBuffer byteBuffer = optionalByteBuffer.get(); + int positionCount = readPositionCount(byteBuffer); + if (positionCount == 0) { + pendingFragment.clear(); + continue; + } + int blockStart = byteBuffer.position(); + int blockSize = byteBuffer.remaining(); + rowCount += positionCount; + if (blockSize <= remainingPayloadSize) { + res.add(byteBuffer); + remainingPayloadSize -= blockSize; + memorySize += blockSize; + pendingFragment.clear(); + } else { + res.add(slice(byteBuffer, blockStart, remainingPayloadSize)); + pendingFragment.reset(byteBuffer, blockStart + remainingPayloadSize); + return ResultChunk.fragment( + res, blockStart + remainingPayloadSize, blockStart == 0 ? blockSize : -1, false); + } + } + return ResultChunk.complete(res, !queryExecution.hasNextResult()); + } + + private static int readPositionCount(ByteBuffer byteBuffer) { + byteBuffer.mark(); + int valueColumnCount = byteBuffer.getInt(); + for (int i = 0; i < valueColumnCount; i++) { + byteBuffer.get(); + } + int positionCount = byteBuffer.getInt(); + byteBuffer.reset(); + return positionCount; + } + + private static ByteBuffer slice(ByteBuffer buffer, int offset, int length) { + ByteBuffer duplicate = buffer.duplicate(); + duplicate.position(offset); + duplicate.limit(offset + length); + return duplicate.slice(); + } + + /** The bounded-payload result of one query result RPC response. */ + public static class ResultChunk { + + private final List queryResult; + private final int nextOffset; + private final int totalLength; + private final boolean finished; + + private ResultChunk( + List queryResult, int nextOffset, int totalLength, boolean finished) { + this.queryResult = queryResult; + this.nextOffset = nextOffset; + this.totalLength = totalLength; + this.finished = finished; + } + + /** Creates a chunk whose last element is a fragment starting at {@code nextOffset}. */ + public static ResultChunk fragment( + List queryResult, int nextOffset, int totalLength, boolean finished) { + return new ResultChunk(queryResult, nextOffset, totalLength, finished); + } + + public static ResultChunk complete(List queryResult, boolean finished) { + return new ResultChunk(queryResult, -1, -1, finished); + } + + public List getQueryResult() { + return queryResult; + } + + public boolean hasNextOffset() { + return nextOffset >= 0; + } + + public int getNextOffset() { + return nextOffset; + } + + public boolean hasTotalLength() { + return totalLength >= 0; + } + + public int getTotalLength() { + return totalLength; + } + + public boolean isFinished() { + return finished; + } + } + public static long[] readTimesFromBuffer(ByteBuffer buffer, int size) { long[] times = new long[size]; for (int i = 0; i < size; i++) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java index b8c53352c09a5..e277418c2a328 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java @@ -298,7 +298,7 @@ private IQueryExecution mockQueryExecution(boolean hasNextResult) throws Excepti IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); Mockito.when(queryExecution.getQueryId()).thenReturn("query"); // the mock carries no buffered batch, only the "is there anything left" flag - Mockito.when(queryExecution.getByteBufferBatchResult()).thenReturn(Optional.empty()); + Mockito.when(queryExecution.getByteBufferBatchResult(false)).thenReturn(Optional.empty()); Mockito.when(queryExecution.hasNextResult()).thenReturn(hasNextResult); return queryExecution; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 95c9b810ce195..f61266b198a44 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -1079,4 +1079,120 @@ public void testForceClose() { Mockito.verify(mockSourceHandleListener, Mockito.timeout(10_0000).times(1)) .onAborted(sourceHandle); } + + @Test + public void testRetainedMemoryIsFreedWhenTheHandedOverBlockIsReleased() throws Exception { + final String queryId = "q0"; + final String localPlanNodeId = "exchange_0"; + final HandleWithPool handleWithPool = + createSourceHandleWithOneFetchedBlock(queryId, localPlanNodeId); + final SourceHandle sourceHandle = handleWithPool.handle; + final MemoryPool memoryPool = handleWithPool.memoryPool; + final String fragmentInstanceId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId( + sourceHandle.getLocalFragmentInstanceId()); + + // The block is taken out of the buffer, but its memory stays reserved in the pool + Assert.assertNotNull(sourceHandle.getSerializedTsBlock(true)); + Assert.assertEquals(0L, sourceHandle.getBufferRetainedSizeInBytes()); + Mockito.verify(memoryPool, Mockito.never()) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + + // The caller releases the memory once the block is fully sent + sourceHandle.freeRetainedMemory(); + Mockito.verify(memoryPool, Mockito.times(1)) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + + // Releasing again is a no-op + sourceHandle.freeRetainedMemory(); + Mockito.verify(memoryPool, Mockito.times(1)) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + } + + @Test + public void testAbortReleasesTheHandedOverMemory() throws Exception { + final String queryId = "q0"; + final String localPlanNodeId = "exchange_0"; + final HandleWithPool handleWithPool = + createSourceHandleWithOneFetchedBlock(queryId, localPlanNodeId); + final SourceHandle sourceHandle = handleWithPool.handle; + final MemoryPool memoryPool = handleWithPool.memoryPool; + final String fragmentInstanceId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId( + sourceHandle.getLocalFragmentInstanceId()); + + Assert.assertNotNull(sourceHandle.getSerializedTsBlock(true)); + Mockito.verify(memoryPool, Mockito.never()) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + + // abort frees the memory that was not released by the caller (fallback path) + sourceHandle.abort(); + Mockito.verify(memoryPool, Mockito.times(1)) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + + // a late release after abort must not free the same memory twice + sourceHandle.freeRetainedMemory(); + Mockito.verify(memoryPool, Mockito.times(1)) + .free(queryId, fragmentInstanceId, localPlanNodeId, MOCK_TSBLOCK_SIZE); + } + + private static final class HandleWithPool { + private final SourceHandle handle; + private final MemoryPool memoryPool; + + private HandleWithPool(SourceHandle handle, MemoryPool memoryPool) { + this.handle = handle; + this.memoryPool = memoryPool; + } + } + + /** Creates a SourceHandle whose single pending block has already been fetched. */ + private HandleWithPool createSourceHandleWithOneFetchedBlock( + String queryId, String localPlanNodeId) throws Exception { + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + + final LocalMemoryManager mockLocalMemoryManager = Mockito.mock(LocalMemoryManager.class); + final MemoryPool mockMemoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(mockMemoryPool); + final IClientManager mockClientManager = + Mockito.mock(IClientManager.class); + final SyncDataNodeMPPDataExchangeServiceClient mockClient = + Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + Mockito.when(mockClientManager.borrowClient(remoteEndpoint)).thenReturn(mockClient); + Mockito.doAnswer( + invocation -> { + TGetDataBlockRequest req = invocation.getArgument(0); + List byteBuffers = new ArrayList<>(); + for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { + byteBuffers.add(ByteBuffer.allocate(0)); + } + return new TGetDataBlockResponse(byteBuffers); + }) + .when(mockClient) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + + final SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + localPlanNodeId, + 0, + mockLocalMemoryManager, + Executors.newSingleThreadExecutor(), + Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE), + Mockito.mock(SourceHandleListener.class), + mockClientManager); + + sourceHandle.isBlocked(); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + Mockito.verify(mockClient, Mockito.timeout(10_000).times(1)) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + Assert.assertTrue(sourceHandle.isBlocked().isDone()); + Assert.assertEquals(MOCK_TSBLOCK_SIZE, sourceHandle.getBufferRetainedSizeInBytes()); + return new HandleWithPool(sourceHandle, mockMemoryPool); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java index cfbcf0e26306c..a21757411cfcd 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java @@ -50,6 +50,7 @@ import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.SeriesScanOptions; import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByKey; import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering; @@ -1808,6 +1809,7 @@ static class FakeQueryExecution implements IQueryExecution { private final long startTime; private final String queryId; private final String sql; + private final PendingResultFragment pendingResultFragment = new PendingResultFragment(() -> {}); FakeQueryExecution(long startTime, String queryId, String sql) { this.startTime = startTime; @@ -1889,10 +1891,15 @@ public long getTimeout() { } @Override - public Optional getByteBufferBatchResult() { + public Optional getByteBufferBatchResult(boolean retainMemory) { return Optional.empty(); } + @Override + public PendingResultFragment getPendingResultFragment() { + return pendingResultFragment; + } + @Override public boolean hasNextResult() { return false; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/QueryDataSetUtilsTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/QueryDataSetUtilsTest.java index 7b65562ce616d..876d792593686 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/QueryDataSetUtilsTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/QueryDataSetUtilsTest.java @@ -21,6 +21,7 @@ import org.apache.iotdb.commons.exception.IoTDBException; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; import org.apache.iotdb.service.rpc.thrift.TSQueryDataSet; import org.apache.tsfile.common.conf.TSFileConfig; @@ -41,6 +42,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Optional; @@ -156,7 +158,7 @@ private IQueryExecution buildQueryExecution() throws IoTDBException, IOException TsBlock tsBlock = buildTsBlock(); Mockito.when(queryExecution.getBatchResult()) .thenReturn(Optional.of(tsBlock), Optional.empty()); - Mockito.when(queryExecution.getByteBufferBatchResult()) + Mockito.when(queryExecution.getByteBufferBatchResult(false)) .thenReturn(Optional.of(new TsBlockSerde().serialize(tsBlock)), Optional.empty()); return queryExecution; } @@ -281,4 +283,195 @@ private void compareTsBlock(Pair, Boolean> res) { new Binary(BINARY_STR, TSFileConfig.STRING_CHARSET), tsBlock.getColumn(5).getBinary(0)); assertTrue(tsBlock.getColumn(5).isNull(1)); } + + @Test + public void testConvertQueryResultByFetchSizeWithSingleOversizedBlock() + throws IoTDBException, IOException { + TsBlock tsBlock = buildTsBlock(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + int totalLength = serialized.remaining(); + // one byte smaller than the block so that the block has to be split into two fragments + int payloadSize = totalLength - 1; + + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Runnable memoryReleaser = Mockito.mock(Runnable.class); + PendingResultFragment pendingFragment = new PendingResultFragment(memoryReleaser); + Mockito.when(queryExecution.getPendingResultFragment()).thenReturn(pendingFragment); + Mockito.when(queryExecution.getByteBufferBatchResult(true)) + .thenReturn(Optional.of(serialized.duplicate()), Optional.empty()); + Mockito.when(queryExecution.hasNextResult()).thenReturn(false); + + QueryDataSetUtils.ResultChunk firstChunk = + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 0, payloadSize); + assertFalse(firstChunk.isFinished()); + assertTrue(firstChunk.hasNextOffset()); + assertTrue(firstChunk.hasTotalLength()); + assertEquals(totalLength, firstChunk.getTotalLength()); + assertEquals(payloadSize, firstChunk.getNextOffset()); + assertEquals(1, firstChunk.getQueryResult().size()); + assertEquals(payloadSize, firstChunk.getQueryResult().get(0).remaining()); + assertTrue(pendingFragment.getBlock() != null); + // the memory of the pending block is still retained + Mockito.verify(memoryReleaser, Mockito.never()).run(); + + QueryDataSetUtils.ResultChunk secondChunk = + QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, 10, firstChunk.getNextOffset(), payloadSize); + assertTrue(secondChunk.isFinished()); + assertFalse(secondChunk.hasNextOffset()); + assertFalse(secondChunk.hasTotalLength()); + assertTrue(pendingFragment.getBlock() == null); + // the memory is released once the fragmented block has been fully sent + Mockito.verify(memoryReleaser, Mockito.times(1)).run(); + + TsBlock result = deserialize(firstChunk.getQueryResult().get(0), secondChunk.getQueryResult()); + assertEquals(2, result.getPositionCount()); + assertEquals(1L, result.getTimeColumn().getLong(0)); + assertEquals(2L, result.getTimeColumn().getLong(1)); + } + + @Test + public void testConvertQueryResultByFetchSizeWithTrailingFragment() + throws IoTDBException, IOException { + TsBlock tsBlock = buildTsBlock(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + int blockSize = serialized.remaining(); + // a small block that fits into the payload, followed by the oversized block + ByteBuffer smallBlock = serialized.duplicate(); + int payloadSize = smallBlock.remaining() + blockSize / 2; + + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Runnable memoryReleaser = Mockito.mock(Runnable.class); + PendingResultFragment pendingFragment = new PendingResultFragment(memoryReleaser); + Mockito.when(queryExecution.getPendingResultFragment()).thenReturn(pendingFragment); + Mockito.when(queryExecution.getByteBufferBatchResult(true)) + .thenReturn(Optional.of(smallBlock), Optional.of(serialized.duplicate()), Optional.empty()); + Mockito.when(queryExecution.hasNextResult()).thenReturn(false); + + QueryDataSetUtils.ResultChunk firstChunk = + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 0, payloadSize); + assertFalse(firstChunk.isFinished()); + assertTrue(firstChunk.hasNextOffset()); + assertEquals(2, firstChunk.getQueryResult().size()); + // the last element is the fragment of the second block + assertEquals(blockSize / 2, firstChunk.getQueryResult().get(1).remaining()); + assertTrue(firstChunk.hasTotalLength()); + assertEquals(blockSize, firstChunk.getTotalLength()); + // the complete block is released at once, the fragmented one is released when it completes + Mockito.verify(memoryReleaser, Mockito.times(1)).run(); + + QueryDataSetUtils.ResultChunk secondChunk = + QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, 10, firstChunk.getNextOffset(), payloadSize); + assertTrue(secondChunk.isFinished()); + assertFalse(secondChunk.hasNextOffset()); + // only the rest of the oversized block remains + assertEquals(1, secondChunk.getQueryResult().size()); + assertEquals(blockSize - blockSize / 2, secondChunk.getQueryResult().get(0).remaining()); + Mockito.verify(memoryReleaser, Mockito.times(2)).run(); + } + + @Test + public void testContinuationCompletionDoesNotFetchNewBlocks() throws IoTDBException, IOException { + TsBlock tsBlock = buildTsBlock(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + int blockSize = serialized.remaining(); + int payloadSize = blockSize / 2; + + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Runnable memoryReleaser = Mockito.mock(Runnable.class); + PendingResultFragment pendingFragment = new PendingResultFragment(memoryReleaser); + Mockito.when(queryExecution.getPendingResultFragment()).thenReturn(pendingFragment); + Mockito.when(queryExecution.getByteBufferBatchResult(true)) + .thenReturn( + Optional.of(serialized.duplicate()), + Optional.of(serialized.duplicate()), + Optional.empty()); + Mockito.when(queryExecution.hasNextResult()).thenReturn(true, false); + + QueryDataSetUtils.ResultChunk firstChunk = + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 0, payloadSize); + assertFalse(firstChunk.isFinished()); + assertEquals(payloadSize, firstChunk.getNextOffset()); + Mockito.verify(memoryReleaser, Mockito.never()).run(); + + // The response that completes the fragmented block must not fetch new blocks, otherwise its + // row count would exceed fetchSize (the tail of the big block is not accounted). + QueryDataSetUtils.ResultChunk continuationChunk = + QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, 10, firstChunk.getNextOffset(), blockSize); + assertFalse(continuationChunk.isFinished()); + assertFalse(continuationChunk.hasNextOffset()); + assertEquals(1, continuationChunk.getQueryResult().size()); + assertEquals(blockSize - payloadSize, continuationChunk.getQueryResult().get(0).remaining()); + assertTrue(pendingFragment.getBlock() == null); + Mockito.verify(memoryReleaser, Mockito.times(1)).run(); + + // New blocks are delivered by the following request only + QueryDataSetUtils.ResultChunk nextChunk = + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 0, blockSize); + assertTrue(nextChunk.isFinished()); + assertFalse(nextChunk.hasNextOffset()); + assertEquals(1, nextChunk.getQueryResult().size()); + assertEquals(blockSize, nextChunk.getQueryResult().get(0).remaining()); + Mockito.verify(memoryReleaser, Mockito.times(2)).run(); + } + + @Test + public void testEmptyBlockReleasesItsMemory() throws IoTDBException, IOException { + TsBlock emptyTsBlock = new TsBlockBuilder(Collections.singletonList(TSDataType.INT64)).build(); + ByteBuffer emptyBlock = new TsBlockSerde().serialize(emptyTsBlock); + TsBlock tsBlock = buildTsBlock(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Runnable memoryReleaser = Mockito.mock(Runnable.class); + Mockito.when(queryExecution.getPendingResultFragment()) + .thenReturn(new PendingResultFragment(memoryReleaser)); + Mockito.when(queryExecution.getByteBufferBatchResult(true)) + .thenReturn(Optional.of(emptyBlock), Optional.of(serialized), Optional.empty()); + Mockito.when(queryExecution.hasNextResult()).thenReturn(false); + + QueryDataSetUtils.ResultChunk chunk = + QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, 10, 0, serialized.remaining()); + assertTrue(chunk.isFinished()); + assertEquals(1, chunk.getQueryResult().size()); + // the empty block and the complete block are both released immediately + Mockito.verify(memoryReleaser, Mockito.times(2)).run(); + } + + @Test(expected = IllegalArgumentException.class) + public void testConvertQueryResultByFetchSizeRejectsMismatchedOffset() throws IoTDBException { + PendingResultFragment pendingFragment = new PendingResultFragment(() -> {}); + pendingFragment.reset(ByteBuffer.wrap(new byte[10]), 5); + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Mockito.when(queryExecution.getPendingResultFragment()).thenReturn(pendingFragment); + + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 3, 8); + } + + @Test(expected = IllegalArgumentException.class) + public void testConvertQueryResultByFetchSizeRejectsOffsetWithoutPendingFragment() + throws IoTDBException { + IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class); + Mockito.when(queryExecution.getPendingResultFragment()) + .thenReturn(new PendingResultFragment(() -> {})); + + QueryDataSetUtils.convertQueryResultByFetchSize(queryExecution, 10, 5, 8); + } + + private TsBlock deserialize(ByteBuffer head, List tail) { + int size = head.remaining(); + for (ByteBuffer buffer : tail) { + size += buffer.remaining(); + } + ByteBuffer assembled = ByteBuffer.allocate(size); + assembled.put(head.duplicate()); + for (ByteBuffer buffer : tail) { + assembled.put(buffer.duplicate()); + } + assembled.flip(); + return new TsBlockSerde().deserialize(assembled); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/ResultChunkingEndToEndTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/ResultChunkingEndToEndTest.java new file mode 100644 index 0000000000000..e7bb00011e68a --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/ResultChunkingEndToEndTest.java @@ -0,0 +1,255 @@ +/* + * 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.utils; + +import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.db.queryengine.plan.execution.PendingResultFragment; +import org.apache.iotdb.rpc.IoTDBRpcDataSet; +import org.apache.iotdb.rpc.RpcUtils; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSCloseOperationReq; +import org.apache.iotdb.service.rpc.thrift.TSFetchResultsReq; +import org.apache.iotdb.service.rpc.thrift.TSFetchResultsResp; + +import org.apache.tsfile.common.conf.TSFileConfig; +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.TsBlockSerde; +import org.apache.tsfile.utils.Binary; +import org.junit.Test; +import org.mockito.Mockito; + +import java.nio.ByteBuffer; +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * End-to-end verification of the fragmented query result protocol: the server side produces + * bounded-payload chunks via {@link QueryDataSetUtils} and the client side ({@code + * IoTDBRpcDataSet}) assembles them back into the original rows. + */ +public class ResultChunkingEndToEndTest { + + private static final String COLUMN = "s1"; + private static final String TEXT_COLUMN = "s2"; + private static final int LARGE_VALUE_SIZE = 512 * 1024; + private static final int MAX_PAYLOAD_SIZE = 256 * 1024; + + @Test(timeout = 60_000) + public void testFragmentedResultRoundTrip() throws Exception { + TsBlock tsBlock = buildTsBlockWithLargeValue(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + int totalLength = serialized.remaining(); + assertTrue( + "the block must be larger than the payload to trigger fragmentation", + totalLength > MAX_PAYLOAD_SIZE); + + ServerSideSimulation server = + new ServerSideSimulation(serialized.duplicate(), tsBlock.getPositionCount(), totalLength); + + IClientRPCService.Iface client = mock(IClientRPCService.Iface.class); + when(client.closeOperation(any(TSCloseOperationReq.class))) + .thenReturn(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenAnswer(invocation -> server.handleFetchResults(invocation.getArgument(0))); + + IoTDBRpcDataSet dataSet = + new IoTDBRpcDataSet( + "", + List.of(COLUMN, TEXT_COLUMN), + List.of(TSDataType.INT64.name(), TSDataType.TEXT.name()), + null, + true, + true, + 1L, + 2L, + client, + 3L, + null, + 100, + 0, + ZoneId.systemDefault(), + RpcUtils.DEFAULT_TIME_FORMAT, + 1_000_000, + true, + new ArrayList<>(Arrays.asList(0, 1))); + + // two rows: [time=1, v=10, text="small"] and [time=2, v=20, text=<512KB>] + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertEquals("small", dataSet.getString(TEXT_COLUMN)); + assertTrue(dataSet.next()); + assertEquals(20L, dataSet.getLong(COLUMN)); + assertEquals(LARGE_VALUE_SIZE, dataSet.getString(TEXT_COLUMN).length()); + assertFalse(dataSet.next()); + + assertTrue( + "the result must have been transferred in several responses", server.rpcCount.get() > 1); + for (int payload : server.payloadSizes) { + assertTrue( + String.format("one response payload %d exceeds the bound %d", payload, MAX_PAYLOAD_SIZE), + payload <= MAX_PAYLOAD_SIZE); + } + // every fragmented block was fully released + assertTrue(server.releaseCount.get() > 0); + assertEquals("the pending fragment must be cleared in the end", null, server.getPendingBlock()); + } + + @Test(timeout = 60_000) + public void testSmallResultUsesOneResponseWithoutFragment() throws Exception { + TsBlockBuilder builder = new TsBlockBuilder(List.of(TSDataType.INT64)); + builder.getTimeColumnBuilder().writeLong(1L); + builder.getColumnBuilder(0).writeLong(10L); + builder.declarePosition(); + TsBlock tsBlock = builder.build(); + ByteBuffer serialized = new TsBlockSerde().serialize(tsBlock); + + ServerSideSimulation server = + new ServerSideSimulation(serialized.duplicate(), tsBlock.getPositionCount(), 0); + + IClientRPCService.Iface client = mock(IClientRPCService.Iface.class); + when(client.closeOperation(any(TSCloseOperationReq.class))) + .thenReturn(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + when(client.fetchResultsV2(any(TSFetchResultsReq.class))) + .thenAnswer(invocation -> server.handleFetchResults(invocation.getArgument(0))); + // the whole result fits into the execute response + IoTDBRpcDataSet dataSet = + new IoTDBRpcDataSet( + "", + List.of(COLUMN), + List.of(TSDataType.INT64.name()), + null, + true, + false, + 1L, + 2L, + client, + 3L, + List.of(serialized.duplicate()), + 100, + 0, + ZoneId.systemDefault(), + RpcUtils.DEFAULT_TIME_FORMAT, + 1_000_000, + true, + new ArrayList<>(List.of(0))); + + assertTrue(dataSet.next()); + assertEquals(10L, dataSet.getLong(COLUMN)); + assertFalse(dataSet.next()); + assertEquals("a small result must not trigger any fetch", 0, server.rpcCount.get()); + } + + /** Mimics ClientRPCServiceImpl#fetchResultsV2 with bounded payloads. */ + private static class ServerSideSimulation { + private final ByteBuffer serialized; + private final int positionCount; + private final int totalLength; + + private final AtomicInteger rpcCount = new AtomicInteger(); + private final AtomicInteger releaseCount = new AtomicInteger(); + private final List payloadSizes = new ArrayList<>(); + private final List remainingBlocks = new ArrayList<>(); + private final PendingResultFragment pendingFragment = + new PendingResultFragment(releaseCount::incrementAndGet); + private final IQueryExecution queryExecution; + + private ServerSideSimulation(ByteBuffer serialized, int positionCount, int totalLength) + throws Exception { + this.serialized = serialized; + this.positionCount = positionCount; + this.totalLength = totalLength; + this.remainingBlocks.add(serialized); + this.queryExecution = Mockito.mock(IQueryExecution.class); + Mockito.when(queryExecution.getPendingResultFragment()).thenReturn(pendingFragment); + Mockito.when(queryExecution.getByteBufferBatchResult(true)) + .thenAnswer( + invocation -> + remainingBlocks.isEmpty() + ? Optional.empty() + : Optional.of(remainingBlocks.remove(0))); + Mockito.when(queryExecution.hasNextResult()) + .thenAnswer( + invocation -> !remainingBlocks.isEmpty() || pendingFragment.getBlock() != null); + } + + private TSFetchResultsResp handleFetchResults(TSFetchResultsReq req) throws Exception { + rpcCount.incrementAndGet(); + QueryDataSetUtils.ResultChunk chunk = + QueryDataSetUtils.convertQueryResultByFetchSize( + queryExecution, + Math.max(positionCount, 1), + req.isSetOffset() ? req.getOffset() : 0, + MAX_PAYLOAD_SIZE); + + int payloadSize = 0; + for (ByteBuffer block : chunk.getQueryResult()) { + payloadSize += block.remaining(); + } + payloadSizes.add(payloadSize); + + TSFetchResultsResp resp = + new TSFetchResultsResp(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS), true, true); + resp.setQueryResult(chunk.getQueryResult()); + resp.setMoreData(!chunk.isFinished()); + if (chunk.hasNextOffset()) { + resp.setOffset(chunk.getNextOffset()); + if (chunk.hasTotalLength()) { + resp.setTotalLength(chunk.getTotalLength()); + } + } + return resp; + } + + private ByteBuffer getPendingBlock() { + return pendingFragment.getBlock(); + } + } + + /** [time=1, v=10, text="small"], [time=2, v=20, text=]. */ + private static TsBlock buildTsBlockWithLargeValue() { + byte[] largeValue = new byte[LARGE_VALUE_SIZE]; + Arrays.fill(largeValue, (byte) 'x'); + + TsBlockBuilder builder = new TsBlockBuilder(Arrays.asList(TSDataType.INT64, TSDataType.TEXT)); + builder.getTimeColumnBuilder().writeLong(1L); + builder.getColumnBuilder(0).writeLong(10L); + builder.getColumnBuilder(1).writeBinary(new Binary("small", TSFileConfig.STRING_CHARSET)); + builder.declarePosition(); + builder.getTimeColumnBuilder().writeLong(2L); + builder.getColumnBuilder(0).writeLong(20L); + builder.getColumnBuilder(1).writeBinary(new Binary(largeValue)); + builder.declarePosition(); + return builder.build(); + } +} diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index a3fe0970318c8..49a23190cf8d5 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -1241,6 +1241,9 @@ mpp_data_exchange_max_pool_size=10 mpp_data_exchange_keep_alive_time_in_ms=1000 # The maximum payload size of one MPP data exchange RPC response. +# It also bounds the payload of one query-result RPC response between client and DataNode: a +# serialized TsBlock larger than this value is transferred in fragments, so that one oversized +# row (e.g. an OBJECT/BLOB column) cannot inflate the RPC buffer of server and client connections. # <=0 use default value # The effective value is between 131072(128KB) and dn_thrift_max_frame_size minus 1024 bytes reserved for the RPC response envelope. # effectiveMode: hot_reload diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift index 7ff12ceb02154..46d2e54554e49 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift @@ -74,6 +74,11 @@ struct TSExecuteStatementResp { 15: optional string database 16: optional bool tableModel 17: optional list columnIndex2TsBlockColumnIndexList + // Start offset of the next fragment. Set only when the last element in queryResult is a + // fragment. Only set for sessions that enabled chunked query results at openSession. + 18: optional i32 offset + // Total serialized length of the TsBlock when the response starts its first fragment. + 19: optional i32 totalLength } enum TSProtocolVersion { @@ -203,6 +208,9 @@ struct TSFetchResultsReq{ 5: required bool isAlign 6: optional i64 timeout 7: optional i64 statementId + // Start offset of the fragment to continue. The client sets this field (including 0 for the + // first request) iff it supports chunked query results. It is optional for rolling upgrade. + 8: optional i32 offset } struct TSFetchResultsResp{ @@ -213,6 +221,10 @@ struct TSFetchResultsResp{ 5: optional TSQueryNonAlignDataSet nonAlignQueryDataSet 6: optional list queryResult 7: optional bool moreData + // Start offset of the next fragment. Set only when the last element in queryResult is a fragment. + 8: optional i32 offset + // Total serialized length of the TsBlock when the response starts its first fragment. + 9: optional i32 totalLength } struct TSFetchMetadataResp{