This is an automated email from the ASF dual-hosted git repository. jt2594838 pushed a commit to branch insert_numeric_row in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 1d6186fface82890ef72cfcaa8279f87f96ca0d8 Author: Tian Jiang <[email protected]> AuthorDate: Wed Sep 16 15:11:10 2026 +0800 tmp --- .../it/session/IoTDBTableModelSessionIT.java | 214 +++++++++++++++++++++ .../org/apache/iotdb/isession/ITableSession.java | 15 ++ .../apache/iotdb/session/i18n/SessionMessages.java | 4 + .../apache/iotdb/session/i18n/SessionMessages.java | 5 + .../apache/iotdb/session/SessionConnection.java | 42 ++++ .../org/apache/iotdb/session/TableSession.java | 12 +- .../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 2 + .../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 4 + .../iotdb/db/protocol/session/IClientSession.java | 53 +++++ .../iotdb/db/protocol/session/SessionManager.java | 2 + .../protocol/thrift/impl/ClientRPCServiceImpl.java | 65 +++++++ 11 files changed, 417 insertions(+), 1 deletion(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java index f9c7ca00444..828ddba1fbd 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java @@ -21,6 +21,7 @@ package org.apache.iotdb.relational.it.session; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.NumericRow; +import org.apache.iotdb.isession.NumericRowPreparedStatement; import org.apache.iotdb.isession.SessionDataSet; import org.apache.iotdb.isession.pool.ITableSessionPool; import org.apache.iotdb.it.env.EnvFactory; @@ -33,16 +34,26 @@ import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.write.record.Tablet; import org.junit.After; +import org.junit.Assume; import org.junit.Before; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; import java.time.LocalDate; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.IntConsumer; import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.showTablesColumnHeaders; import static org.junit.Assert.assertEquals; @@ -182,6 +193,209 @@ public class IoTDBTableModelSessionIT { } } + @Test + public void testPreparedNumericRowAndTabletWritePerformance() throws Exception { + Assume.assumeTrue( + "Performance benchmark disabled; set iotdb.numeric.row.insert.perf.enabled=true to run", + Boolean.getBoolean("iotdb.numeric.row.insert.perf.enabled")); + final int concurrency = 16; + // Compare sustained single-row writes after warming both paths; verify all acknowledged rows. + final int warmupRounds = 2; + final int measuredRounds = 6; + final int fieldCount = 99; + try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { + session.executeNonQueryStatement("CREATE DATABASE numeric_perf"); + session.executeNonQueryStatement("USE numeric_perf"); + StringBuilder ddl = new StringBuilder("CREATE TABLE prepared_perf (tag STRING TAG"); + for (int i = 0; i < fieldCount; i++) { + ddl.append(", f").append(i).append(" DOUBLE FIELD"); + } + ddl.append(")"); + session.executeNonQueryStatement(ddl.toString()); + session.executeNonQueryStatement(ddl.toString().replace("prepared_perf", "tablet_perf")); + } + + List<String> names = new ArrayList<>(); + List<TSDataType> types = new ArrayList<>(); + List<ColumnCategory> categories = new ArrayList<>(); + names.add("tag"); + types.add(TSDataType.STRING); + categories.add(ColumnCategory.TAG); + for (int i = 0; i < fieldCount; i++) { + names.add("f" + i); + types.add(TSDataType.DOUBLE); + categories.add(ColumnCategory.FIELD); + } + ITableSession[] sessions = new ITableSession[concurrency]; + NumericRowPreparedStatement[] prepared = new NumericRowPreparedStatement[concurrency]; + Tablet[] tablets = new Tablet[concurrency]; + String[][] tags = new String[concurrency][]; + long[][] numericValues = new long[concurrency][fieldCount]; + long[] preparedRows = new long[concurrency]; + long[] tabletRows = new long[concurrency]; + ExecutorService executor = Executors.newFixedThreadPool(concurrency); + try { + // Connections, USE, prepare and reusable input allocation are outside every timed round. + for (int i = 0; i < concurrency; i++) { + sessions[i] = EnvFactory.getEnv().getTableSessionConnection(); + sessions[i].executeNonQueryStatement("USE numeric_perf"); + prepared[i] = + sessions[i].prepareInsertNumericRow("prepared_perf", names, types, categories); + tablets[i] = new Tablet("tablet_perf", names, types, categories, 1); + tags[i] = new String[] {"device-" + i}; + Arrays.fill(numericValues[i], Double.doubleToRawLongBits(1.0)); + } + IntConsumer preparedWrite = + threadId -> { + try { + sessions[threadId].executePreparedInsertNumericRow( + prepared[threadId], + preparedRows[threadId], + numericValues[threadId], + tags[threadId]); + preparedRows[threadId]++; + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + IntConsumer tabletWrite = + threadId -> { + try { + Tablet tablet = tablets[threadId]; + tablet.addTimestamp(0, tabletRows[threadId]); + tablet.addValue("tag", 0, tags[threadId][0]); + for (int i = 0; i < fieldCount; i++) { + tablet.addValue(names.get(i + 1), 0, 1.0D); + } + sessions[threadId].insert(tablet); + tablet.reset(); + tabletRows[threadId]++; + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + double[][] rates = new double[2][measuredRounds]; + long[] totalRows = new long[2]; + long[] totalNanos = new long[2]; + for (int round = 0; round < warmupRounds + measuredRounds; round++) { + boolean warmup = round < warmupRounds; + int seconds = warmup ? 3 : 5; + // Alternate AB/BA to distribute shared JVM warmup and background storage work. + for (int position = 0; position < 2; position++) { + int path = (round + position) % 2; + WritePerformance result = + runConcurrentWrites( + executor, concurrency, seconds, path == 0 ? preparedWrite : tabletWrite); + System.out.printf( + "%s round %d position %d %s: %d rows, %.3f s, %.1f rows/s%n", + warmup ? "warmup" : "measured", + warmup ? round + 1 : round - warmupRounds + 1, + position + 1, + path == 0 ? "prepared numeric row" : "insertTablet batch=1", + result.rows(), + result.nanos() / 1_000_000_000.0, + result.rowsPerSecond()); + if (!warmup) { + rates[path][round - warmupRounds] = result.rowsPerSecond(); + totalRows[path] += result.rows(); + totalNanos[path] += result.nanos(); + } + } + } + for (int path = 0; path < 2; path++) { + Arrays.sort(rates[path]); + System.out.printf( + "summary %s: median %.1f rows/s, aggregate %.1f rows/s, min %.1f, max %.1f%n", + path == 0 ? "prepared numeric row" : "insertTablet batch=1", + (rates[path][measuredRounds / 2 - 1] + rates[path][measuredRounds / 2]) / 2, + totalRows[path] * 1_000_000_000.0 / totalNanos[path], + rates[path][0], + rates[path][measuredRounds - 1]); + } + // Per-device monotonic timestamps span warmup and measurement, so no rows overwrite others. + for (int path = 0; path < 2; path++) { + String table = path == 0 ? "prepared_perf" : "tablet_perf"; + long expectedRows = Arrays.stream(path == 0 ? preparedRows : tabletRows).sum(); + try (SessionDataSet dataSet = + sessions[0].executeQueryStatement("SELECT count(*) FROM " + table)) { + assertTrue(dataSet.hasNext()); + assertEquals(expectedRows, dataSet.next().getFields().get(0).getLongV()); + assertFalse(dataSet.hasNext()); + } + System.out.printf("verified %s: %d rows%n", table, expectedRows); + } + } finally { + executor.shutdownNow(); + try { + assertTrue(executor.awaitTermination(1, TimeUnit.MINUTES)); + } finally { + // Closing each session also releases its server-side prepared metadata. + Exception closeFailure = null; + for (ITableSession session : sessions) { + if (session != null) { + try { + session.close(); + } catch (Exception e) { + if (closeFailure == null) closeFailure = e; + else closeFailure.addSuppressed(e); + } + } + } + if (closeFailure != null) throw closeFailure; + } + } + } + + private record WritePerformance(long rows, long nanos) { + double rowsPerSecond() { + return rows * 1_000_000_000.0 / nanos; + } + } + + private static WritePerformance runConcurrentWrites( + ExecutorService executor, int concurrency, int seconds, IntConsumer write) throws Exception { + CountDownLatch ready = new CountDownLatch(concurrency); + CountDownLatch start = new CountDownLatch(1); + AtomicLong deadline = new AtomicLong(); + List<Future<WritePerformance>> futures = new ArrayList<>(); + try { + for (int i = 0; i < concurrency; i++) { + final int threadId = i; + futures.add( + executor.submit( + () -> { + ready.countDown(); + start.await(); + long rows = 0; + while (!Thread.currentThread().isInterrupted() + && System.nanoTime() < deadline.get()) { + write.accept(threadId); + rows++; + } + return new WritePerformance(rows, System.nanoTime()); + })); + } + assertTrue(ready.await(30, TimeUnit.SECONDS)); + long startNanos = System.nanoTime(); + deadline.set(startNanos + TimeUnit.SECONDS.toNanos(seconds)); + start.countDown(); + long rows = 0; + long endNanos = startNanos; + for (Future<WritePerformance> future : futures) { + WritePerformance result = future.get(seconds + 60, TimeUnit.SECONDS); + rows += result.rows(); + endNanos = Math.max(endNanos, result.nanos()); + } + // Include completion of in-flight RPCs, but exclude main-thread collection/printing time. + return new WritePerformance(rows, endNanos - startNanos); + } finally { + start.countDown(); + for (Future<WritePerformance> future : futures) { + if (!future.isDone()) future.cancel(true); + } + } + } + private void assertNumericTagRows(ITableSession session) throws Exception { String[] devices = {"设备甲😀", "设备乙", "", null}; try (SessionDataSet dataSet = diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java index af524f9a853..d14abab1238 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java @@ -22,8 +22,12 @@ package org.apache.iotdb.isession; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.write.record.Tablet; +import java.util.List; + /** * This interface defines a session for interacting with IoTDB tables. It supports operations such * as data insertion, executing queries, and closing the session. Implementations of this interface @@ -58,6 +62,17 @@ public interface ITableSession extends AutoCloseable { */ void insert(NumericRow row) throws StatementExecutionException, IoTDBConnectionException; + NumericRowPreparedStatement prepareInsertNumericRow( + String tableName, + List<String> columnNames, + List<TSDataType> dataTypes, + List<ColumnCategory> columnCategories) + throws StatementExecutionException, IoTDBConnectionException; + + void executePreparedInsertNumericRow( + NumericRowPreparedStatement prepared, long timestamp, long[] values, String[] tagValues) + throws StatementExecutionException, IoTDBConnectionException; + /** * Executes a non-query SQL statement, such as a DDL or DML command. * diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index 3eb19971a92..173260e4a07 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -290,4 +290,8 @@ public final class SessionMessages { public static final String EXCEPTION_NUMERIC_ROW_ARGUMENTS_MUST_BE_NON_NULL_AND_COLUMN_NAMES_TYPES_CATEGORIES_AND_VALUES_MUST_HAVE_EQUAL_LENGTHS_AB102B01 = "Numeric row arguments must be non-null and column names, types, categories and values must have equal lengths."; public static final String EXCEPTION_NUMERIC_ROWS_SUPPORT_ONLY_FIELD_COLUMNS_DA7A8EE7 = "Numeric rows support only FIELD columns."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_BELONGS_TO_ANOTHER_CONNECTION_PREPARE_IT_AGAIN_055CC7BC = "Prepared numeric row is closed or belongs to another connection; prepare it again."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_PREPARE_RETURNED_NO_STATUS_6FF7BC24 = "Prepared numeric row prepare returned no status"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_COUNT_MISMATCH_9ACD3A0C = "Prepared numeric row value count mismatch"; + public static final String EXCEPTION_UNSUPPORTED_NUMERIC_TYPE_06DD3044 = "Unsupported numeric type"; } diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 22f3bea7821..b45704d8738 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -288,4 +288,9 @@ public final class SessionMessages { public static final String EXCEPTION_NUMERIC_ROW_ARGUMENTS_MUST_BE_NON_NULL_AND_COLUMN_NAMES_TYPES_CATEGORIES_AND_VALUES_MUST_HAVE_EQUAL_LENGTHS_AB102B01 = "数值行参数不能为 null，列名、类型、类别和值的数量必须相同。"; public static final String EXCEPTION_NUMERIC_ROWS_SUPPORT_ONLY_FIELD_COLUMNS_DA7A8EE7 = "数值行仅支持 FIELD 列。"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_BELONGS_TO_ANOTHER_CONNECTION_PREPARE_IT_AGAIN_055CC7BC = + "预备数值行已关闭或属于其他连接，请重新准备"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_PREPARE_RETURNED_NO_STATUS_6FF7BC24 = "预备数值行准备请求未返回状态"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_COUNT_MISMATCH_9ACD3A0C = "预备数值行的值数量不匹配"; + public static final String EXCEPTION_UNSUPPORTED_NUMERIC_TYPE_06DD3044 = "不支持的数值类型"; } 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 2e144448b58..e01c5dd8a20 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 @@ -44,6 +44,7 @@ import org.apache.iotdb.service.rpc.thrift.TSCreateSchemaTemplateReq; import org.apache.iotdb.service.rpc.thrift.TSCreateTimeseriesReq; import org.apache.iotdb.service.rpc.thrift.TSDeleteDataReq; import org.apache.iotdb.service.rpc.thrift.TSDropSchemaTemplateReq; +import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedNumericRowReq; import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementReq; import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp; import org.apache.iotdb.service.rpc.thrift.TSFastLastDataQueryForOneDeviceReq; @@ -59,6 +60,8 @@ import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq; import org.apache.iotdb.service.rpc.thrift.TSLastDataQueryReq; import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowReq; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowResp; import org.apache.iotdb.service.rpc.thrift.TSPruneSchemaTemplateReq; import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateReq; import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateResp; @@ -79,6 +82,7 @@ import org.apache.tsfile.utils.Preconditions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.nio.ByteBuffer; import java.security.SecureRandom; import java.time.ZoneId; import java.util.ArrayList; @@ -759,6 +763,44 @@ public class SessionConnection { callWithRetryAndVerifyWithRedirection(() -> insertRecordInternal(request)); } + protected long prepareNumericRow(TSPrepareNumericRowReq request) + throws IoTDBConnectionException, StatementExecutionException { + request.setSessionId(sessionId); + try { + TSPrepareNumericRowResp resp = client.prepareNumericRow(request); + if (resp.getStatus() == null) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_PREPARED_NUMERIC_ROW_PREPARE_RETURNED_NO_STATUS_6FF7BC24); + } + RpcUtils.verifySuccess(resp.getStatus()); + return resp.getPreparedId(); + } catch (TException e) { + throw new IoTDBConnectionException(e); + } + } + + protected void executePreparedNumericRow( + long id, long timestamp, ByteBuffer values, ByteBuffer tags) + throws IoTDBConnectionException, StatementExecutionException { + TSExecutePreparedNumericRowReq request = + new TSExecutePreparedNumericRowReq(sessionId, id, timestamp, values, tags); + try { + RpcUtils.verifySuccess(client.executePreparedNumericRow(request)); + } catch (TException e) { + throw new IoTDBConnectionException(e); + } + } + + protected void closePreparedNumericRow(long id) throws IoTDBConnectionException { + try { + RpcUtils.verifySuccess(client.closePreparedNumericRow(sessionId, id)); + } catch (StatementExecutionException e) { + throw new IoTDBConnectionException(e); + } catch (TException e) { + throw new IoTDBConnectionException(e); + } + } + private TSStatus insertRecordInternal(TSInsertRecordReq request) throws TException { request.setSessionId(sessionId); return client.insertRecord(request); diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java index 38cc6d72cea..f7154b3c5f7 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java @@ -21,17 +21,27 @@ package org.apache.iotdb.session; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.NumericRow; +import org.apache.iotdb.isession.NumericRowPreparedStatement; import org.apache.iotdb.isession.SessionDataSet; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowReq; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.record.Tablet; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; + public class TableSession implements ITableSession { private final Session session; - TableSession(Session session) { + public TableSession(Session session) { this.session = session; } diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index 12f4e374ca7..d4ce7f97ba8 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1489,4 +1489,6 @@ public final class DataNodeMiscMessages { LOG_TABLE_QUERY_DEVICE_ENTRY_BATCH_SIZE_IN_BYTES_ARG_EXCEEDS_DN_THRIFT_MAX_FRAME_SIZE_ARG_USING_ARG_AS_THE_EFFECTIVE_VALUE_2AE1BEDA = "table_query_device_entry_batch_size_in_bytes (%d) exceeds the maximum RPC payload (dn_thrift_max_frame_size %d minus 1024 bytes); using %d as the effective value"; + public static final String EXCEPTION_INVALID_PREPARED_NUMERIC_ROW_METADATA_OR_PAYLOAD_AECCCEAA = "Invalid prepared numeric row metadata or payload."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A = "Prepared numeric row is closed or does not belong to this session."; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index d809fea1992..2d19c74be8b 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1469,4 +1469,8 @@ public final class DataNodeMiscMessages { LOG_TABLE_QUERY_DEVICE_ENTRY_BATCH_SIZE_IN_BYTES_ARG_EXCEEDS_DN_THRIFT_MAX_FRAME_SIZE_ARG_USING_ARG_AS_THE_EFFECTIVE_VALUE_2AE1BEDA = "table_query_device_entry_batch_size_in_bytes（%d）超过最大 RPC payload（dn_thrift_max_frame_size %d 减去 1024 字节），将使用 %d 作为生效值"; + public static final String EXCEPTION_INVALID_PREPARED_NUMERIC_ROW_METADATA_OR_PAYLOAD_AECCCEAA = + "预备数值行元数据或数据载荷无效"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A = + "预备数值行已关闭或不属于当前会话"; } 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 bac4b15e342..0283b226bdf 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 @@ -26,6 +26,7 @@ import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.utils.CommonDateTimeUtils; import org.apache.iotdb.commons.utils.PathUtils; import org.apache.iotdb.db.i18n.DataNodeMiscMessages; +import org.apache.iotdb.db.queryengine.plan.execution.config.session.PreparedStatementMemoryManager; import org.apache.iotdb.rpc.subscription.annotation.TableModel; import org.apache.iotdb.service.rpc.thrift.TSConnectionInfo; import org.apache.iotdb.service.rpc.thrift.TSConnectionType; @@ -33,12 +34,64 @@ import org.apache.iotdb.service.rpc.thrift.TSConnectionType; import javax.annotation.Nullable; import java.time.ZoneId; +import java.util.HashMap; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.TimeZone; +import java.util.concurrent.atomic.AtomicLong; public abstract class IClientSession { + private static final AtomicLong NUMERIC_ROW_IDS = new AtomicLong(); + private Map<Long, PreparedNumericRowInfo> preparedNumericRows; + private boolean numericRowsClosed; + + public synchronized long addPreparedNumericRow(PreparedNumericRowInfo info) { + if (numericRowsClosed) { + throw new SemanticException( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A); + } + long preparedId = NUMERIC_ROW_IDS.incrementAndGet(); + if (preparedNumericRows == null) { + preparedNumericRows = new HashMap<>(); + } + PreparedStatementMemoryManager.getInstance() + .allocate(Long.toString(preparedId), info.ramBytesUsed()); + preparedNumericRows.put(preparedId, info); + return preparedId; + } + + public synchronized PreparedNumericRowInfo getPreparedNumericRow(long preparedId) { + PreparedNumericRowInfo info = + preparedNumericRows == null ? null : preparedNumericRows.get(preparedId); + if (info == null) { + throw new SemanticException( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A); + } + return info; + } + + public synchronized void removePreparedNumericRow(long preparedId) { + PreparedNumericRowInfo info = + preparedNumericRows == null ? null : preparedNumericRows.remove(preparedId); + if (info != null) { + PreparedStatementMemoryManager.getInstance().release(info.ramBytesUsed()); + } + } + + public synchronized void releasePreparedNumericRows() { + numericRowsClosed = true; + if (preparedNumericRows != null) { + for (PreparedNumericRowInfo info : preparedNumericRows.values()) { + PreparedStatementMemoryManager.getInstance().release(info.ramBytesUsed()); + } + preparedNumericRows = null; + } + } + private long id; private ClientVersion clientVersion; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 21dbffb606e..fa1da5d8ea5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -260,6 +260,8 @@ public class SessionManager implements SessionManagerMBean { } } + session.releasePreparedNumericRows(); + // Release PreparedStatement memory resources try { PreparedStatementMemoryManager.getInstance().releaseAllForSession(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 99d40b3572b..d48aef4cbf6 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 @@ -69,6 +69,7 @@ import org.apache.iotdb.db.protocol.client.ConfigNodeClient; import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager; import org.apache.iotdb.db.protocol.client.ConfigNodeInfo; import org.apache.iotdb.db.protocol.session.IClientSession; +import org.apache.iotdb.db.protocol.session.PreparedNumericRowInfo; import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.protocol.thrift.OperationType; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; @@ -170,6 +171,7 @@ import org.apache.iotdb.service.rpc.thrift.TSDeallocatePreparedReq; import org.apache.iotdb.service.rpc.thrift.TSDeleteDataReq; import org.apache.iotdb.service.rpc.thrift.TSDropSchemaTemplateReq; import org.apache.iotdb.service.rpc.thrift.TSExecuteBatchStatementReq; +import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedNumericRowReq; import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedReq; import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementReq; import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp; @@ -191,6 +193,8 @@ import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq; import org.apache.iotdb.service.rpc.thrift.TSLastDataQueryReq; import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowReq; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowResp; import org.apache.iotdb.service.rpc.thrift.TSPrepareReq; import org.apache.iotdb.service.rpc.thrift.TSPrepareResp; import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; @@ -2582,6 +2586,67 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler { } } + @Override + public TSPrepareNumericRowResp prepareNumericRow(TSPrepareNumericRowReq req) { + IClientSession session = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime(); + if (!SESSION_MANAGER.checkLogin(session)) { + return new TSPrepareNumericRowResp(getNotLoggedInStatus()); + } + try { + if (req.getSessionId() != session.getId()) { + return new TSPrepareNumericRowResp(RpcUtils.getStatus(TSStatusCode.NOT_LOGIN)); + } + PreparedNumericRowInfo info = new PreparedNumericRowInfo(session.getDatabaseName(), req); + TSPrepareNumericRowResp response = new TSPrepareNumericRowResp(); + response.setStatus(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + response.setPreparedId(session.addPreparedNumericRow(info)); + return response; + } catch (Exception e) { + return new TSPrepareNumericRowResp( + onQueryException( + e, OperationType.INSERT_RECORD.getName(), TSStatusCode.EXECUTE_STATEMENT_ERROR)); + } + } + + @Override + public TSStatus executePreparedNumericRow(TSExecutePreparedNumericRowReq req) { + IClientSession session = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime(); + if (!SESSION_MANAGER.checkLogin(session)) return getNotLoggedInStatus(); + try { + if (req.getSessionId() != session.getId()) return RpcUtils.getStatus(TSStatusCode.NOT_LOGIN); + InsertBaseStatement statement = + session + .getPreparedNumericRow(req.getPreparedId()) + .bind( + req.getTimestamp(), + ByteBuffer.wrap(req.getTagValues()), + ByteBuffer.wrap(req.getValues())); + long queryId = SESSION_MANAGER.requestQueryId(); + return COORDINATOR.executeForTableModel( + statement, + relationSqlParser, + session, + queryId, + SESSION_MANAGER.getSessionInfo(session), + "", + metadata, + config.getConnectionTimeoutInMS()) + .status; + } catch (Exception e) { + return onQueryException( + e, OperationType.INSERT_RECORD.getName(), TSStatusCode.EXECUTE_STATEMENT_ERROR); + } + } + + @Override + public TSStatus closePreparedNumericRow(long sessionId, long preparedId) { + IClientSession session = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime(); + if (!SESSION_MANAGER.checkLogin(session) || session.getId() != sessionId) + return getNotLoggedInStatus(); + session.removePreparedNumericRow(preparedId); + return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + } + @Override public TSStatus insertRecord(TSInsertRecordReq req) { long t1 = System.nanoTime();
