You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ji...@apache.org on 2026/10/12 03:39:01 UTC
(iotdb) 18/20: fix multiple issues
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 2794bfc6adcb444d3973f1a00e859384c6f4d3ec
Author: Tian Jiang <jt...@163.com>
AuthorDate: Wed Sep 30 11:09:20 2026 +0800
fix multiple issues
---
.../manual/basic/IoTDBPipeNumericRowIT.java | 449 +++++++++++++++++++++
.../it/session/IoTDBTableModelSessionIT.java | 206 ++++++++--
.../it/session/WritePerformanceIoTracker.java | 215 ++++++++++
.../apache/iotdb/session/i18n/SessionMessages.java | 1 +
.../apache/iotdb/session/i18n/SessionMessages.java | 1 +
.../org/apache/iotdb/session/TableSession.java | 29 ++
.../iotdb/session/pool/TableSessionWrapper.java | 35 +-
.../session/PreparedNumericRowMetadataTest.java | 218 ++++++++++
.../session/pool/PreparedNumericRowPoolTest.java | 222 ++++++++++
.../statement/PipeConvertedInsertRowStatement.java | 35 +-
...leStatementDataTypeConvertExecutionVisitor.java | 10 +
.../protocol/thrift/impl/ClientRPCServiceImpl.java | 8 +
.../PipeConvertedInsertRowStatementTest.java | 146 +++++++
...ClientRPCServiceImplPreparedNumericRowTest.java | 195 +++++++++
.../plan/parser/NumericTableRowTest.java | 61 +++
15 files changed, 1784 insertions(+), 47 deletions(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/basic/IoTDBPipeNumericRowIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/basic/IoTDBPipeNumericRowIT.java
new file mode 100644
index 00000000000..5211f290771
--- /dev/null
+++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/basic/IoTDBPipeNumericRowIT.java
@@ -0,0 +1,449 @@
+/*
+ * 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.pipe.it.dual.tablemodel.manual.basic;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
+import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
+import org.apache.iotdb.commons.cluster.NodeStatus;
+import org.apache.iotdb.commons.cluster.RegionRoleType;
+import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
+import org.apache.iotdb.confignode.rpc.thrift.TRegionInfo;
+import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowRegionReq;
+import org.apache.iotdb.confignode.rpc.thrift.TShowRegionResp;
+import org.apache.iotdb.consensus.ConsensusFactory;
+import org.apache.iotdb.isession.ITableSession;
+import org.apache.iotdb.isession.NumericRow;
+import org.apache.iotdb.isession.NumericRowPreparedStatement;
+import org.apache.iotdb.isession.SessionConfig;
+import org.apache.iotdb.it.env.MultiEnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.MultiClusterIT2DualTableManualBasic;
+import org.apache.iotdb.itbase.env.BaseEnv;
+import org.apache.iotdb.pipe.it.dual.tablemodel.manual.AbstractPipeTableModelDualManualIT;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.session.TableSessionBuilder;
+
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.utils.BitMap;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.time.LocalDate;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(IoTDBTestRunner.class)
+@Category(MultiClusterIT2DualTableManualBasic.class)
+public class IoTDBPipeNumericRowIT extends AbstractPipeTableModelDualManualIT {
+ private static final String DATABASE = "numeric_sync";
+ private static final List<String> NAMES =
+ Arrays.asList("d", "device", "f", "l", "region", "i", "b", "event_date", "event_time");
+ private static final List<TSDataType> TYPES =
+ Arrays.asList(
+ TSDataType.DOUBLE,
+ TSDataType.STRING,
+ TSDataType.FLOAT,
+ TSDataType.INT64,
+ TSDataType.STRING,
+ TSDataType.INT32,
+ TSDataType.BOOLEAN,
+ TSDataType.DATE,
+ TSDataType.TIMESTAMP);
+ private static final List<ColumnCategory> CATEGORIES =
+ Arrays.asList(
+ ColumnCategory.FIELD,
+ ColumnCategory.TAG,
+ ColumnCategory.FIELD,
+ ColumnCategory.FIELD,
+ ColumnCategory.TAG,
+ ColumnCategory.FIELD,
+ ColumnCategory.FIELD,
+ ColumnCategory.FIELD,
+ ColumnCategory.FIELD);
+
+ @Override
+ @Before
+ public void setUp() {
+ MultiEnvFactory.createEnv(2);
+ senderEnv = MultiEnvFactory.getEnv(0);
+ receiverEnv = MultiEnvFactory.getEnv(1);
+ setupConfig();
+ senderEnv
+ .getConfig()
+ .getCommonConfig()
+ .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS)
+ .setDataRegionGroupExtensionPolicy("CUSTOM")
+ .setDefaultDataRegionGroupNumPerDatabase(1)
+ .setDataReplicationFactor(3)
+ .setSchemaReplicationFactor(3);
+ senderEnv
+ .getConfig()
+ .getDataNodeConfig()
+ .setMetricReporterType(Collections.singletonList("PROMETHEUS"));
+ receiverEnv
+ .getConfig()
+ .getCommonConfig()
+ .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS);
+ senderEnv.initClusterEnvironment(1, 3);
+ receiverEnv.initClusterEnvironment(1, 1);
+ }
+
+ @Test
+ public void testRealtimeWithoutBatch() throws Exception {
+ // Unflushed NumericRow and prepared rows must arrive through individual InsertNode requests.
+ testRealtime(false);
+ }
+
+ @Test
+ public void testRealtimeWithBatch() throws Exception {
+ // Batch transfer must preserve all numeric types, FIELD null bits, and null/empty STRING TAGs.
+ testRealtime(true);
+ }
+
+ private void testRealtime(boolean batch) throws Exception {
+ createTable();
+ createPipe(batch, false);
+ writeRows(0, 32);
+ assertRowsEventually(receiverEnv, 0, 32);
+ // The destination must remain readable after its own MemTable has been flushed.
+ try (ITableSession session = receiverEnv.getTableSessionConnection()) {
+ session.executeNonQueryStatement("FLUSH");
+ }
+ assertRowsEventually(receiverEnv, 0, 32);
+ }
+
+ @Test
+ public void testHistoricalTsFile() throws Exception {
+ // Start the pipe after flushing so the historical TsFile path also carries numeric-row data.
+ createTable();
+ writeRows(0, 32);
+ try (ITableSession session = senderEnv.getTableSessionConnection()) {
+ session.executeNonQueryStatement("FLUSH");
+ }
+ createPipe(true, true);
+ assertRowsEventually(receiverEnv, 0, 32);
+ }
+
+ @Test
+ public void testIoTConsensusReplicaAndRestart() throws Exception {
+ // Pin writes to the actual DataRegion leader, then stop that leader before querying followers.
+ // Merely stopping a client entry node can still leave all reads on the original leader.
+ createTable();
+ writeRows(0, 16);
+ int leader = findDataRegionLeader();
+ awaitReplication(leader, 16);
+ int offlineFollower = (leader + 1) % 3;
+ senderEnv.shutdownDataNode(offlineFollower);
+ try {
+ awaitNodeStatus(offlineFollower, NodeStatus.Unknown);
+ writeRows(16, 32, leader);
+ } finally {
+ senderEnv.startDataNode(offlineFollower);
+ }
+ awaitNodeStatus(offlineFollower, NodeStatus.Running);
+ // Running only confirms node health; IoTConsensus catch-up must finish before losing its
+ // sender.
+ awaitReplication(leader, 32);
+ // Resolve the leader again after restart instead of assuming its assignment is unchanged.
+ int stoppedLeader = findDataRegionLeader();
+ assertEquals(
+ "The primary that accepted the second write batch must be stopped", leader, stoppedLeader);
+ senderEnv.shutdownDataNode(stoppedLeader);
+ try {
+ assertFalse(senderEnv.getDataNodeWrapper(stoppedLeader).isAlive());
+ awaitNodeStatus(stoppedLeader, NodeStatus.Unknown);
+ for (int follower = 0; follower < 3; follower++) {
+ if (follower == stoppedLeader) {
+ continue;
+ }
+ assertFalse(senderEnv.getDataNodeWrapper(stoppedLeader).isAlive());
+ assertRowsEventually(senderEnv, follower, 32);
+ assertFalse(senderEnv.getDataNodeWrapper(stoppedLeader).isAlive());
+ System.out.printf(
+ "NumericRow IoTConsensus: stopped primary port=%d, query follower port=%d, verified rows=32%n",
+ senderEnv.getDataNodeWrapper(stoppedLeader).getPort(),
+ senderEnv.getDataNodeWrapper(follower).getPort());
+ }
+ } finally {
+ senderEnv.startDataNode(stoppedLeader);
+ }
+ }
+
+ private void awaitReplication(int leader, int expectedWrites) {
+ String metricsUrl =
+ "http://"
+ + senderEnv.getDataNodeWrapper(leader).getIp()
+ + ":"
+ + senderEnv.getDataNodeWrapper(leader).getMetricPort()
+ + "/metrics";
+ await()
+ .atMost(90, TimeUnit.SECONDS)
+ .pollInterval(1, TimeUnit.SECONDS)
+ .untilAsserted(
+ () -> {
+ String metrics = senderEnv.getUrlContent(metricsUrl, null);
+ assertNotNull(metrics);
+ // Require the latest write index too: an old zero-lag sample is not catch-up proof.
+ // This database has exactly one DataRegion and each numeric row adds one write.
+ double searchIndex = -1;
+ double safeIndex = -1;
+ for (String line : metrics.split("\\R")) {
+ if (!line.startsWith("iot_consensus{")) {
+ continue;
+ }
+ double value = Double.parseDouble(line.substring(line.lastIndexOf(' ') + 1));
+ if (line.contains("type=\"searchIndex\"")) {
+ searchIndex = value;
+ } else if (line.contains("type=\"safeIndex\"")) {
+ safeIndex = value;
+ }
+ }
+ assertTrue(
+ "Missing or stale searchIndex at " + metricsUrl, searchIndex >= expectedWrites);
+ assertEquals("Replication is incomplete at " + metricsUrl, searchIndex, safeIndex, 0);
+ });
+ }
+
+ private void awaitNodeStatus(int node, NodeStatus expected) {
+ // BaseEnv.ensureNodeStatus uses EnvFactory, which cannot be used with MultiCluster
+ // environments.
+ await()
+ .atMost(90, TimeUnit.SECONDS)
+ .pollInterval(1, TimeUnit.SECONDS)
+ .untilAsserted(
+ () -> {
+ try (SyncConfigNodeIServiceClient client =
+ (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) {
+ TShowClusterResp response = client.showCluster();
+ assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode());
+ List<TDataNodeLocation> matches =
+ response.getDataNodeList().stream()
+ .filter(
+ location ->
+ location.getClientRpcEndPoint().getPort()
+ == senderEnv.getDataNodeWrapper(node).getPort())
+ .toList();
+ assertEquals(1, matches.size());
+ assertEquals(
+ expected.getStatus(),
+ response.getNodeStatus().get(matches.get(0).getDataNodeId()));
+ }
+ });
+ }
+
+ private int findDataRegionLeader() throws Exception {
+ try (SyncConfigNodeIServiceClient client =
+ (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) {
+ TShowRegionResp response =
+ client.showRegion(
+ new TShowRegionReq()
+ .setConsensusGroupType(TConsensusGroupType.DataRegion)
+ .setDatabases(Collections.singletonList(DATABASE))
+ .setIsTableModel(true));
+ assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode());
+ // Exactly one three-replica region ensures stopping this leader covers all inserted rows.
+ List<TRegionInfo> replicas = response.getRegionInfoList();
+ assertEquals(3, replicas.size());
+ for (TRegionInfo replica : replicas) {
+ assertEquals(replicas.get(0).getConsensusGroupId(), replica.getConsensusGroupId());
+ }
+ List<TRegionInfo> leaders =
+ replicas.stream()
+ .filter(replica -> RegionRoleType.Leader.getRoleType().equals(replica.getRoleType()))
+ .toList();
+ assertEquals(1, leaders.size());
+ TRegionInfo leader = leaders.get(0);
+ for (int node = 0; node < 3; node++) {
+ if (senderEnv.getDataNodeWrapper(node).getPort() == leader.getClientRpcPort()) {
+ System.out.printf(
+ "NumericRow IoTConsensus: region=%d, primary DataNodeId=%d, port=%d%n",
+ leader.getConsensusGroupId().getId(),
+ leader.getDataNodeId(),
+ leader.getClientRpcPort());
+ return node;
+ }
+ }
+ throw new AssertionError("No wrapper for DataRegion leader " + leader);
+ }
+ }
+
+ private void createTable() throws Exception {
+ try (ITableSession session = senderEnv.getTableSessionConnection()) {
+ session.executeNonQueryStatement("CREATE DATABASE " + DATABASE);
+ session.executeNonQueryStatement("USE " + DATABASE);
+ session.executeNonQueryStatement(
+ "CREATE TABLE readings(region STRING TAG, device STRING TAG, missing STRING TAG, "
+ + "d DOUBLE FIELD, f FLOAT FIELD, l INT64 FIELD, i INT32 FIELD, b BOOLEAN FIELD, "
+ + "event_date DATE FIELD, event_time TIMESTAMP FIELD)");
+ }
+ }
+
+ private void createPipe(boolean batch, boolean history) throws Exception {
+ Map<String, String> source = new HashMap<>();
+ source.put("capture.table", "true");
+ source.put("capture.tree", "false");
+ source.put("__system.sql-dialect", "table");
+ source.put("user", "root");
+ source.put("database-name", DATABASE);
+ source.put("table-name", "readings");
+ source.put("source.realtime.mode", "log");
+ source.put("source.history.enable", Boolean.toString(history));
+ Map<String, String> sink = new HashMap<>();
+ sink.put("sink", "iotdb-thrift-sink");
+ sink.put("sink.ip", receiverEnv.getDataNodeWrapper(0).getIp());
+ sink.put("sink.port", Integer.toString(receiverEnv.getDataNodeWrapper(0).getPort()));
+ sink.put("sink.batch.enable", Boolean.toString(batch));
+ sink.put("sink.batch.max-delay-seconds", "1");
+ sink.put("sink.format", history ? "tsfile" : "tablet");
+ try (SyncConfigNodeIServiceClient client =
+ (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) {
+ assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ client
+ .createPipe(
+ new TCreatePipeReq("numeric_pipe", sink)
+ .setExtractorAttributes(source)
+ .setProcessorAttributes(Collections.emptyMap()))
+ .getCode());
+ assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(), client.startPipe("numeric_pipe").getCode());
+ }
+ }
+
+ private void writeRows(int start, int end) throws Exception {
+ writeRows(start, end, 0);
+ }
+
+ private void writeRows(int start, int end, int node) throws Exception {
+ // Pin the client to the selected endpoint; its index is not necessarily the DataRegion leader.
+ try (ITableSession session =
+ new TableSessionBuilder()
+ .nodeUrls(
+ Collections.singletonList(
+ senderEnv.getDataNodeWrapper(node).getIpAndPortString()))
+ .username(SessionConfig.DEFAULT_USER)
+ .password(SessionConfig.DEFAULT_PASSWORD)
+ .database(DATABASE)
+ .enableRedirection(false)
+ .build();
+ NumericRowPreparedStatement prepared =
+ session.prepareInsertNumericRow("readings", NAMES, TYPES, CATEGORIES)) {
+ for (int row = start; row < end; row++) {
+ long[] values = {
+ Double.doubleToRawLongBits(row + 0.25),
+ Float.floatToRawIntBits(-row - 0.5f),
+ Long.MIN_VALUE + row,
+ Integer.MIN_VALUE + row,
+ row % 2,
+ 20260924,
+ 1000L + row
+ };
+ String[] tags = {device(row), "region"};
+ BitMap nulls = new BitMap(7);
+ if (row % 3 == 0) {
+ nulls.mark(row % 7);
+ }
+ if (row % 2 == 0) {
+ session.executePreparedInsertNumericRow(prepared, row, values, tags, nulls);
+ } else {
+ session.insert(
+ new NumericRow("readings", row, NAMES, TYPES, CATEGORIES, values, tags, nulls));
+ }
+ }
+ }
+ }
+
+ private static String device(int row) {
+ return row % 3 == 0 ? null : row % 3 == 1 ? "" : "device-\u8bbe\u5907";
+ }
+
+ private void assertRowsEventually(BaseEnv env, int node, int count) {
+ await()
+ // Restart and leader failover can temporarily reject JDBC connections or query RPCs.
+ .ignoreExceptionsInstanceOf(SQLException.class)
+ .atMost(90, TimeUnit.SECONDS)
+ .pollInterval(1, TimeUnit.SECONDS)
+ .untilAsserted(
+ () -> {
+ try (Connection connection =
+ env.getConnection(
+ env.getDataNodeWrapper(node),
+ SessionConfig.DEFAULT_USER,
+ SessionConfig.DEFAULT_PASSWORD,
+ "table");
+ Statement statement = connection.createStatement()) {
+ statement.execute("USE " + DATABASE);
+ try (ResultSet result =
+ statement.executeQuery(
+ "SELECT time,region,device,missing,d,f,l,i,b,event_date,event_time FROM readings ORDER BY time")) {
+ for (int row = 0; row < count; row++) {
+ assertTrue("Missing row " + row, result.next());
+ assertEquals(row, result.getLong(1));
+ assertEquals("region", result.getString(2));
+ assertEquals(device(row), result.getString(3));
+ assertNull(result.getObject(4));
+ for (int field = 0; field < 7; field++) {
+ int column = field + 5;
+ if (row % 3 == 0 && field == row % 7) {
+ assertNull(result.getObject(column));
+ continue;
+ }
+ assertNotNull(result.getObject(column));
+ switch (field) {
+ case 0 -> assertEquals(row + 0.25, result.getDouble(column), 0);
+ case 1 -> assertEquals(-row - 0.5f, result.getFloat(column), 0);
+ case 2 -> assertEquals(Long.MIN_VALUE + row, result.getLong(column));
+ case 3 -> assertEquals(Integer.MIN_VALUE + row, result.getInt(column));
+ case 4 -> assertEquals(row % 2 != 0, result.getBoolean(column));
+ case 5 ->
+ assertEquals(
+ LocalDate.of(2026, 9, 24), result.getDate(column).toLocalDate());
+ case 6 -> assertEquals(1000L + row, result.getLong(column));
+ default -> throw new AssertionError(field);
+ }
+ }
+ }
+ assertFalse(result.next());
+ }
+ }
+ });
+ }
+}
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 3b0313ec006..e196ff1edf0 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
@@ -72,6 +72,16 @@ public class IoTDBTableModelSessionIT {
@Before
public void setUp() throws Exception {
+ if (Boolean.getBoolean("iotdb.numeric.row.insert.perf.enabled")
+ && Boolean.getBoolean("iotdb.numeric.row.insert.perf.io.enabled")) {
+ // Isolate per-round storage accounting from compaction of earlier rounds.
+ EnvFactory.getEnv()
+ .getConfig()
+ .getCommonConfig()
+ .setEnableSeqSpaceCompaction(false)
+ .setEnableUnseqSpaceCompaction(false)
+ .setEnableCrossSpaceCompaction(false);
+ }
EnvFactory.getEnv().initClusterEnvironment();
}
@@ -431,10 +441,11 @@ public class IoTDBTableModelSessionIT {
@Test
public void testPreparedNumericRowAndTabletWritePerformance() throws Exception {
+ // Compare synchronous single-row paths and optionally account for their Linux TCP/storage I/O.
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;
+ final int concurrency = Integer.getInteger("iotdb.numeric.row.insert.perf.concurrency", 16);
// Compare sustained single-row writes after warming all three paths; verify all acknowledged
// rows.
final int warmupRounds = 2;
@@ -497,30 +508,21 @@ public class IoTDBTableModelSessionIT {
pathRows[0][threadId],
numericValues[threadId],
tags[threadId]);
- pathRows[0][threadId]++;
} catch (Exception e) {
throw new RuntimeException(e);
}
};
- IntConsumer tabletWrite =
+ IntConsumer tabletPrepare =
threadId -> {
- try {
- Tablet tablet = tablets[threadId];
- tablet.addTimestamp(0, pathRows[2][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();
- pathRows[2][threadId]++;
- } catch (Exception e) {
- throw new RuntimeException(e);
+ Tablet tablet = tablets[threadId];
+ tablet.addTimestamp(0, pathRows[2][threadId]);
+ tablet.addValue("tag", 0, tags[threadId][0]);
+ for (int i = 0; i < fieldCount; i++) {
+ tablet.addValue(names.get(i + 1), 0, 1.0D);
}
};
- IntConsumer rowWrite =
- threadId -> {
- try {
+ IntConsumer rowPrepare =
+ threadId ->
rows[threadId] =
new NumericRow(
"row_perf",
@@ -530,51 +532,112 @@ public class IoTDBTableModelSessionIT {
categories,
numericValues[threadId],
tags[threadId]);
- sessions[threadId].insert(rows[threadId]);
- pathRows[1][threadId]++;
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- };
- IntConsumer[] paths = {preparedWrite, rowWrite, tabletWrite};
+ IntConsumer[] paths = {
+ preparedWrite,
+ threadId -> {
+ try {
+ sessions[threadId].insert(rows[threadId]);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ },
+ threadId -> {
+ try {
+ sessions[threadId].insert(tablets[threadId]);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+ };
+ IntConsumer[] prepareWrites = {null, rowPrepare, tabletPrepare};
+ IntConsumer[] completeWrites = {
+ threadId -> pathRows[0][threadId]++,
+ threadId -> pathRows[1][threadId]++,
+ threadId -> {
+ tablets[threadId].reset();
+ pathRows[2][threadId]++;
+ }
+ };
String[] pathNames = {"prepared numeric row", "insert NumericRow", "insertTablet batch=1"};
System.out.printf("configuration: concurrency=%d, tabletBatch=1%n", concurrency);
+ WritePerformanceIoTracker ioTracker = null;
+ if (Boolean.getBoolean("iotdb.numeric.row.insert.perf.io.enabled")) {
+ assertEquals(
+ "I/O measurement requires one local DataNode",
+ 1,
+ EnvFactory.getEnv().getDataNodeWrapperList().size());
+ assertEquals(
+ "I/O measurement requires one local ConfigNode",
+ 1,
+ EnvFactory.getEnv().getConfigNodeWrapperList().size());
+ ioTracker =
+ new WritePerformanceIoTracker(
+ EnvFactory.getEnv().getDataNodeWrapperList().get(0).getPid(),
+ EnvFactory.getEnv().getConfigNodeWrapperList().get(0).getPid(),
+ EnvFactory.getEnv().getDataNodeWrapperList().get(0).getPort());
+ }
double[][] rates = new double[paths.length][measuredRounds];
long[] totalRows = new long[paths.length];
long[] totalNanos = new long[paths.length];
+ LongSamples[] requestLatencies = new LongSamples[paths.length];
+ for (int path = 0; path < paths.length; path++) {
+ requestLatencies[path] = new LongSamples();
+ }
for (int round = 0; round < warmupRounds + measuredRounds; round++) {
boolean warmup = round < warmupRounds;
int seconds = warmup ? 3 : 5;
// Rotate all three paths so each occupies every position during measured rounds.
for (int position = 0; position < paths.length; position++) {
int path = (round + position) % paths.length;
+ WritePerformanceIoTracker.Snapshot ioBefore =
+ ioTracker == null ? null : ioTracker.beforeRound(sessions[0]);
WritePerformance result =
- runConcurrentWrites(executor, concurrency, seconds, paths[path]);
+ runConcurrentWrites(
+ executor,
+ concurrency,
+ seconds,
+ prepareWrites[path],
+ paths[path],
+ completeWrites[path]);
+ if (ioTracker != null) {
+ System.out.printf(
+ "io %s round %d position %d %s: rows=%d, %s%n",
+ warmup ? "warmup" : "measured",
+ warmup ? round + 1 : round - warmupRounds + 1,
+ position + 1,
+ pathNames[path],
+ result.rows(),
+ ioTracker.afterRound(ioBefore, sessions[0]));
+ }
System.out.printf(
- "%s round %d position %d %s: %d rows, %.3f s, %.1f rows/s%n",
+ "%s round %d position %d %s: %d rows, %.3f s, %.1f rows/s, API latency %s%n",
warmup ? "warmup" : "measured",
warmup ? round + 1 : round - warmupRounds + 1,
position + 1,
pathNames[path],
result.rows(),
result.nanos() / 1_000_000_000.0,
- result.rowsPerSecond());
+ result.rowsPerSecond(),
+ summarizeRequestLatencies(result.requestNanos()));
if (!warmup) {
rates[path][round - warmupRounds] = result.rowsPerSecond();
totalRows[path] += result.rows();
totalNanos[path] += result.nanos();
+ requestLatencies[path].appendAll(result.requestNanos());
}
}
}
for (int path = 0; path < paths.length; path++) {
Arrays.sort(rates[path]);
System.out.printf(
- "summary %s: median %.1f rows/s, aggregate %.1f rows/s, min %.1f, max %.1f%n",
+ "summary %s: median %.1f rows/s, aggregate %.1f rows/s, min %.1f, max %.1f, "
+ + "measured API latency %s%n",
pathNames[path],
(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]);
+ rates[path][measuredRounds - 1],
+ summarizeRequestLatencies(requestLatencies[path].toArray()));
}
// Per-device monotonic timestamps span warmup and measurement, so no rows overwrite others.
for (int path = 0; path < paths.length; path++) {
@@ -619,14 +682,80 @@ public class IoTDBTableModelSessionIT {
}
}
- private record WritePerformance(long rows, long nanos) {
+ private record WritePerformance(long rows, long nanos, long[] requestNanos) {
double rowsPerSecond() {
return rows * 1_000_000_000.0 / nanos;
}
}
+ private record RequestLatency(
+ long count, double averageNanos, long p50Nanos, long p95Nanos, long p99Nanos) {
+ @Override
+ public String toString() {
+ return String.format(
+ "n=%d, avg=%.3f us, p50=%.3f us, p95=%.3f us, p99=%.3f us",
+ count,
+ averageNanos / 1_000.0,
+ p50Nanos / 1_000.0,
+ p95Nanos / 1_000.0,
+ p99Nanos / 1_000.0);
+ }
+ }
+
+ private static class LongSamples {
+ private long[] values = new long[1024];
+ private int size;
+
+ void add(long value) {
+ if (size == values.length) {
+ values = Arrays.copyOf(values, values.length * 2);
+ }
+ values[size++] = value;
+ }
+
+ void appendAll(long[] samples) {
+ while (values.length - size < samples.length) {
+ values = Arrays.copyOf(values, values.length * 2);
+ }
+ System.arraycopy(samples, 0, values, size, samples.length);
+ size += samples.length;
+ }
+
+ long[] toArray() {
+ return Arrays.copyOf(values, size);
+ }
+ }
+
+ private static RequestLatency summarizeRequestLatencies(long[] requestNanos) {
+ if (requestNanos.length == 0) {
+ return new RequestLatency(0, 0, 0, 0, 0);
+ }
+ long[] sorted = Arrays.copyOf(requestNanos, requestNanos.length);
+ Arrays.sort(sorted);
+ double sum = 0;
+ for (long latency : sorted) {
+ sum += latency;
+ }
+ return new RequestLatency(
+ sorted.length,
+ sum / sorted.length,
+ percentile(sorted, 0.50),
+ percentile(sorted, 0.95),
+ percentile(sorted, 0.99));
+ }
+
+ private static long percentile(long[] sorted, double quantile) {
+ return sorted[Math.max(0, (int) Math.ceil(quantile * sorted.length) - 1)];
+ }
+
private static WritePerformance runConcurrentWrites(
- ExecutorService executor, int concurrency, int seconds, IntConsumer write) throws Exception {
+ ExecutorService executor,
+ int concurrency,
+ int seconds,
+ IntConsumer prepare,
+ IntConsumer write,
+ IntConsumer complete)
+ throws Exception {
CountDownLatch ready = new CountDownLatch(concurrency);
CountDownLatch start = new CountDownLatch(1);
AtomicLong deadline = new AtomicLong();
@@ -640,12 +769,19 @@ public class IoTDBTableModelSessionIT {
ready.countDown();
start.await();
long rows = 0;
+ LongSamples requestNanos = new LongSamples();
while (!Thread.currentThread().isInterrupted()
&& System.nanoTime() < deadline.get()) {
+ if (prepare != null) {
+ prepare.accept(threadId);
+ }
+ long requestStart = System.nanoTime();
write.accept(threadId);
+ requestNanos.add(System.nanoTime() - requestStart);
+ complete.accept(threadId);
rows++;
}
- return new WritePerformance(rows, System.nanoTime());
+ return new WritePerformance(rows, System.nanoTime(), requestNanos.toArray());
}));
}
assertTrue(ready.await(30, TimeUnit.SECONDS));
@@ -654,13 +790,15 @@ public class IoTDBTableModelSessionIT {
start.countDown();
long rows = 0;
long endNanos = startNanos;
+ LongSamples requestNanos = new LongSamples();
for (Future<WritePerformance> future : futures) {
WritePerformance result = future.get(seconds + 60, TimeUnit.SECONDS);
rows += result.rows();
endNanos = Math.max(endNanos, result.nanos());
+ requestNanos.appendAll(result.requestNanos());
}
// Include completion of in-flight RPCs, but exclude main-thread collection/printing time.
- return new WritePerformance(rows, endNanos - startNanos);
+ return new WritePerformance(rows, endNanos - startNanos, requestNanos.toArray());
} finally {
start.countDown();
for (Future<WritePerformance> future : futures) {
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/WritePerformanceIoTracker.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/WritePerformanceIoTracker.java
new file mode 100644
index 00000000000..bab6f0ad710
--- /dev/null
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/WritePerformanceIoTracker.java
@@ -0,0 +1,215 @@
+/*
+ * 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.relational.it.session;
+
+import org.apache.iotdb.isession.ITableSession;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.TimeUnit;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/** Linux-only, opt-in counters sampled outside the single-row benchmark's timed write loops. */
+final class WritePerformanceIoTracker {
+ private static final Pattern TCP_COUNTER =
+ Pattern.compile("\\b(bytes_received|bytes_sent|bytes_retrans):(\\d+)\\b");
+
+ private final long dataNodePid;
+ private final long configNodePid;
+ private final int rpcPort;
+
+ WritePerformanceIoTracker(long dataNodePid, long configNodePid, int rpcPort) throws IOException {
+ assertTrue("I/O measurement requires Linux /proc", Files.isDirectory(Path.of("/proc")));
+ assertTrue("DataNode PID unavailable", dataNodePid > 0);
+ assertTrue("ConfigNode PID unavailable", configNodePid > 0);
+ this.dataNodePid = dataNodePid;
+ this.configNodePid = configNodePid;
+ this.rpcPort = rpcPort;
+ // Check readability early; unsupported counters must fail rather than appear as zero I/O.
+ readDisk(dataNodePid);
+ readDisk(configNodePid);
+ System.out.printf(
+ "io configuration: dataNodePid=%d, configNodePid=%d, rpcPort=%d, "
+ + "network=server TCP payload counters, disk=/proc/pid/io, "
+ + "flushBeforeAndAfter=true, compaction=false%n",
+ dataNodePid, configNodePid, rpcPort);
+ }
+
+ Snapshot beforeRound(ITableSession session) throws Exception {
+ // Drain previous MemTables/WAL before the baseline; the following flush is charged to this
+ // round.
+ session.executeNonQueryStatement("FLUSH");
+ return new Snapshot(readNetwork(), readDisk(dataNodePid), readDisk(configNodePid));
+ }
+
+ Result afterRound(Snapshot before, ITableSession session) throws Exception {
+ // Snapshot TCP before FLUSH so SQL/control traffic is excluded from write-request bytes.
+ Map<String, Traffic> network = readNetwork();
+ assertEquals(
+ "RPC connections changed during I/O measurement",
+ before.network.keySet(),
+ network.keySet());
+ Traffic traffic = new Traffic(0, 0, 0);
+ for (Map.Entry<String, Traffic> entry : network.entrySet()) {
+ traffic = traffic.plus(entry.getValue().minus(before.network.get(entry.getKey())));
+ }
+ Disk activeDataNode = readDisk(dataNodePid).minus(before.dataNode);
+ long flushStart = System.nanoTime();
+ session.executeNonQueryStatement("FLUSH");
+ long flushNanos = System.nanoTime() - flushStart;
+ Disk dataNode = readDisk(dataNodePid).minus(before.dataNode);
+ Disk configNode = readDisk(configNodePid).minus(before.configNode);
+ return new Result(traffic, activeDataNode, dataNode, configNode, flushNanos);
+ }
+
+ private Map<String, Traffic> readNetwork() throws Exception {
+ // Count only server endpoints on this test DataNode's RPC port. Loopback RX/TX are not doubled.
+ Process process =
+ new ProcessBuilder("ss", "-tinHO", "state", "established", "( sport = :" + rpcPort + " )")
+ .redirectErrorStream(true)
+ .start();
+ CompletableFuture<String> output =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ return new String(process.getInputStream().readAllBytes(), StandardCharsets.UTF_8);
+ } catch (IOException e) {
+ throw new CompletionException(e);
+ }
+ });
+ try {
+ assertTrue("ss did not finish", process.waitFor(10, TimeUnit.SECONDS));
+ String text = output.get(10, TimeUnit.SECONDS);
+ assertEquals(text, 0, process.exitValue());
+ Map<String, Traffic> result = new HashMap<>();
+ for (String line : text.split("\\R")) {
+ if (line.isBlank()) continue;
+ String[] fields = line.trim().split("\\s+", 5);
+ assertTrue("Unexpected ss output: " + line, fields.length == 5);
+ Map<String, Long> counters = new HashMap<>();
+ Matcher matcher = TCP_COUNTER.matcher(fields[4]);
+ while (matcher.find()) counters.put(matcher.group(1), Long.parseLong(matcher.group(2)));
+ assertTrue(
+ "TCP byte counters unavailable: " + line,
+ counters.containsKey("bytes_received") && counters.containsKey("bytes_sent"));
+ String connection = fields[2] + "->" + fields[3];
+ assertFalse("Duplicate server socket: " + connection, result.containsKey(connection));
+ // ss omits zero-valued bytes_retrans; the received/sent counters above are mandatory.
+ result.put(
+ connection,
+ new Traffic(
+ counters.get("bytes_received"),
+ counters.get("bytes_sent"),
+ counters.getOrDefault("bytes_retrans", 0L)));
+ }
+ assertFalse("No established test RPC connections", result.isEmpty());
+ return result;
+ } finally {
+ process.destroyForcibly();
+ output.cancel(true);
+ }
+ }
+
+ private static Disk readDisk(long pid) throws IOException {
+ Map<String, Long> counters = new HashMap<>();
+ for (String line : Files.readAllLines(Path.of("/proc", Long.toString(pid), "io"))) {
+ String[] fields = line.split(":\\s*", 2);
+ counters.put(fields[0], Long.parseLong(fields[1].trim()));
+ }
+ assertTrue(
+ "Process storage byte counters unavailable for PID " + pid,
+ counters.containsKey("read_bytes")
+ && counters.containsKey("write_bytes")
+ && counters.containsKey("cancelled_write_bytes"));
+ // These are kernel-accounted process storage bytes, not rchar/wchar (which include socket I/O),
+ // nor physical NAND bytes. Preserve cancelled writes separately instead of hiding them.
+ return new Disk(
+ counters.get("read_bytes"),
+ counters.get("write_bytes"),
+ counters.get("cancelled_write_bytes"));
+ }
+
+ record Snapshot(Map<String, Traffic> network, Disk dataNode, Disk configNode) {}
+
+ private record Traffic(long received, long sent, long retransmitted) {
+ Traffic minus(Traffic previous) {
+ assertTrue(
+ "TCP counters decreased",
+ received >= previous.received
+ && sent >= previous.sent
+ && retransmitted >= previous.retransmitted);
+ return new Traffic(
+ received - previous.received,
+ sent - previous.sent,
+ retransmitted - previous.retransmitted);
+ }
+
+ Traffic plus(Traffic other) {
+ return new Traffic(
+ received + other.received, sent + other.sent, retransmitted + other.retransmitted);
+ }
+ }
+
+ private record Disk(long read, long written, long cancelled) {
+ Disk minus(Disk previous) {
+ assertTrue(
+ "Process I/O counters decreased",
+ read >= previous.read && written >= previous.written && cancelled >= previous.cancelled);
+ return new Disk(
+ read - previous.read, written - previous.written, cancelled - previous.cancelled);
+ }
+ }
+
+ record Result(
+ Traffic traffic, Disk activeDataNode, Disk dataNode, Disk configNode, long flushNanos) {
+ @Override
+ public String toString() {
+ return String.format(
+ "requestBytes=%d, responseBytes=%d, responseRetransmitBytes=%d, "
+ + "dnActiveReadBytes=%d, dnActiveWriteBytes=%d, dnActiveCancelledBytes=%d, "
+ + "dnReadBytes=%d, dnWriteBytes=%d, dnCancelledBytes=%d, "
+ + "cnReadBytes=%d, cnWriteBytes=%d, cnCancelledBytes=%d, flushNanos=%d",
+ traffic.received,
+ traffic.sent,
+ traffic.retransmitted,
+ activeDataNode.read,
+ activeDataNode.written,
+ activeDataNode.cancelled,
+ dataNode.read,
+ dataNode.written,
+ dataNode.cancelled,
+ configNode.read,
+ configNode.written,
+ configNode.cancelled,
+ flushNanos);
+ }
+ }
+}
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 a1bdf45262e..2bcac9214ce 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
@@ -296,4 +296,5 @@ public final class SessionMessages {
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";
+ public static final String EXCEPTION_TABLE_SESSION_IS_CLOSED_77FE359B = "Table session is closed.";
}
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 260b000e61a..aa9c1055fd5 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
@@ -295,4 +295,5 @@ public final class SessionMessages {
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 = "不支持的数值类型";
+ public static final String EXCEPTION_TABLE_SESSION_IS_CLOSED_77FE359B = "表会话已关闭。";
}
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 9d3e6d8b9bb..c453596c1ee 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
@@ -67,6 +67,35 @@ public class TableSession implements ITableSession {
List<TSDataType> dataTypes,
List<ColumnCategory> columnCategories)
throws StatementExecutionException, IoTDBConnectionException {
+ // Validate before regrouping, which would otherwise discard unsupported categories or
+ // surplus metadata entries.
+ if (tableName == null
+ || columnNames == null
+ || dataTypes == null
+ || columnCategories == null
+ || columnNames.size() != dataTypes.size()
+ || columnNames.size() != columnCategories.size()) {
+ throw new IllegalArgumentException(
+ SessionMessages
+ .EXCEPTION_NUMERIC_ROW_ARGUMENTS_MUST_BE_NON_NULL_COLUMN_METADATA_LENGTHS_AND_PER_CATEGORY_VALUE_COUNTS_MUST_MATCH_EB12B233);
+ }
+ for (int i = 0; i < columnNames.size(); i++) {
+ if (columnNames.get(i) == null
+ || dataTypes.get(i) == null
+ || columnCategories.get(i) == null) {
+ throw new IllegalArgumentException(
+ SessionMessages
+ .EXCEPTION_NUMERIC_ROW_ARGUMENTS_MUST_BE_NON_NULL_COLUMN_METADATA_LENGTHS_AND_PER_CATEGORY_VALUE_COUNTS_MUST_MATCH_EB12B233);
+ }
+ if ((columnCategories.get(i) != ColumnCategory.TAG
+ && columnCategories.get(i) != ColumnCategory.FIELD)
+ || (columnCategories.get(i) == ColumnCategory.TAG
+ && dataTypes.get(i) != TSDataType.STRING)) {
+ throw new IllegalArgumentException(
+ SessionMessages
+ .EXCEPTION_NUMERIC_ROWS_SUPPORT_NUMERIC_FIELD_COLUMNS_AND_STRING_TAG_COLUMNS_ONLY_87159FFA);
+ }
+ }
List<String> names = new ArrayList<>();
List<TSDataType> types = new ArrayList<>();
List<ColumnCategory> categories = new ArrayList<>();
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java
index 7c950949d0d..7684beba9ab 100644
--- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java
@@ -80,6 +80,7 @@ public class TableSessionWrapper implements ITableSession {
@Override
public void insert(NumericRow row) throws StatementExecutionException, IoTDBConnectionException {
+ checkOpen();
try {
session.insertRelationalRecord(row);
} catch (IoTDBConnectionException e) {
@@ -97,8 +98,14 @@ public class TableSessionWrapper implements ITableSession {
List<TSDataType> dataTypes,
List<ColumnCategory> columnCategories)
throws StatementExecutionException, IoTDBConnectionException {
- return tableSession.prepareInsertNumericRow(
- tableName, columnNames, dataTypes, columnCategories);
+ checkOpen();
+ try {
+ return tableSession.prepareInsertNumericRow(
+ tableName, columnNames, dataTypes, columnCategories);
+ } catch (IoTDBConnectionException e) {
+ invalidateSession();
+ throw e;
+ }
}
@Override
@@ -116,17 +123,32 @@ public class TableSessionWrapper implements ITableSession {
String[] tagValues,
BitMap nullBitmap)
throws StatementExecutionException, IoTDBConnectionException {
+ checkOpen();
try {
tableSession.executePreparedInsertNumericRow(
prepared, timestamp, values, tagValues, nullBitmap);
} catch (IoTDBConnectionException e) {
- sessionPool.cleanSessionAndMayThrowConnectionException(session);
- closed.set(true);
- session = null;
+ invalidateSession();
throw e;
}
}
+ private void checkOpen() throws IoTDBConnectionException {
+ // The cached delegate still references the physical session after it returns to the pool.
+ if (closed.get() || session == null) {
+ throw new IoTDBConnectionException(
+ SessionMessages.EXCEPTION_TABLE_SESSION_IS_CLOSED_77FE359B);
+ }
+ }
+
+ private void invalidateSession() {
+ Session failedSession = session;
+ // Invalidate before cleanup so close() cannot return a failed connection to the pool.
+ closed.set(true);
+ session = null;
+ sessionPool.cleanSessionAndMayThrowConnectionException(failedSession);
+ }
+
@Override
public SessionDataSet executeQueryStatement(String sql)
throws StatementExecutionException, IoTDBConnectionException {
@@ -168,6 +190,9 @@ public class TableSessionWrapper implements ITableSession {
@Override
public void close() throws IoTDBConnectionException {
+ if (closed.get() || session == null) {
+ return;
+ }
if (!Objects.equals(session.getSqlDialect(), sessionPool.sqlDialect)) {
try {
session.executeNonQueryStatement("set sql_dialect=" + sessionPool.sqlDialect);
diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowMetadataTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowMetadataTest.java
new file mode 100644
index 00000000000..b47b71a961d
--- /dev/null
+++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowMetadataTest.java
@@ -0,0 +1,218 @@
+/*
+ * 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.session;
+
+import org.apache.iotdb.isession.NumericRowPreparedStatement;
+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.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+public class PreparedNumericRowMetadataTest {
+ private static final String INVALID_ARGUMENTS =
+ SessionMessages
+ .EXCEPTION_NUMERIC_ROW_ARGUMENTS_MUST_BE_NON_NULL_COLUMN_METADATA_LENGTHS_AND_PER_CATEGORY_VALUE_COUNTS_MUST_MATCH_EB12B233;
+ private static final String INVALID_CATEGORIES =
+ SessionMessages
+ .EXCEPTION_NUMERIC_ROWS_SUPPORT_NUMERIC_FIELD_COLUMNS_AND_STRING_TAG_COLUMNS_ONLY_87159FFA;
+
+ @Test
+ public void testNullMetadataRejectedBeforeRpc() throws Exception {
+ assertInvalid(
+ null,
+ List.of("v"),
+ List.of(TSDataType.DOUBLE),
+ List.of(ColumnCategory.FIELD),
+ INVALID_ARGUMENTS);
+ assertInvalid(
+ "t", null, List.of(TSDataType.DOUBLE), List.of(ColumnCategory.FIELD), INVALID_ARGUMENTS);
+ assertInvalid("t", List.of("v"), null, List.of(ColumnCategory.FIELD), INVALID_ARGUMENTS);
+ assertInvalid("t", List.of("v"), List.of(TSDataType.DOUBLE), null, INVALID_ARGUMENTS);
+ }
+
+ @Test
+ public void testNullMetadataElementsRejectedBeforeRpc() throws Exception {
+ assertInvalid(
+ "t",
+ Arrays.asList("tag", null),
+ List.of(TSDataType.STRING, TSDataType.DOUBLE),
+ List.of(ColumnCategory.TAG, ColumnCategory.FIELD),
+ INVALID_ARGUMENTS);
+ assertInvalid(
+ "t",
+ List.of("tag", "v"),
+ Arrays.asList(TSDataType.STRING, null),
+ List.of(ColumnCategory.TAG, ColumnCategory.FIELD),
+ INVALID_ARGUMENTS);
+ assertInvalid(
+ "t",
+ List.of("tag", "v"),
+ List.of(TSDataType.STRING, TSDataType.DOUBLE),
+ Arrays.asList(ColumnCategory.TAG, null),
+ INVALID_ARGUMENTS);
+ }
+
+ @Test
+ public void testMetadataLengthMismatchRejectedBeforeRpc() throws Exception {
+ final List<String> names = List.of("a", "b");
+ final List<TSDataType> types = List.of(TSDataType.DOUBLE, TSDataType.DOUBLE);
+ final List<ColumnCategory> categories = List.of(ColumnCategory.FIELD, ColumnCategory.FIELD);
+ assertInvalid("t", names.subList(0, 1), types, categories, INVALID_ARGUMENTS);
+ assertInvalid("t", names, types.subList(0, 1), categories, INVALID_ARGUMENTS);
+ assertInvalid("t", names, types, categories.subList(0, 1), INVALID_ARGUMENTS);
+ assertInvalid(
+ "t",
+ names,
+ List.of(TSDataType.DOUBLE, TSDataType.DOUBLE, TSDataType.DOUBLE),
+ categories,
+ INVALID_ARGUMENTS);
+ assertInvalid(
+ "t",
+ names,
+ types,
+ List.of(ColumnCategory.FIELD, ColumnCategory.FIELD, ColumnCategory.FIELD),
+ INVALID_ARGUMENTS);
+ }
+
+ @Test
+ public void testAttributeColumnRejectedBeforeRpc() throws Exception {
+ assertInvalid(
+ "t",
+ List.of("tag", "attr", "v"),
+ List.of(TSDataType.STRING, TSDataType.STRING, TSDataType.DOUBLE),
+ List.of(ColumnCategory.TAG, ColumnCategory.ATTRIBUTE, ColumnCategory.FIELD),
+ INVALID_CATEGORIES);
+ }
+
+ @Test
+ public void testNonStringTagRejectedBeforeRpc() throws Exception {
+ for (TSDataType type : List.of(TSDataType.INT64, TSDataType.TEXT)) {
+ assertInvalid(
+ "t",
+ List.of("tag", "v"),
+ List.of(type, TSDataType.DOUBLE),
+ List.of(ColumnCategory.TAG, ColumnCategory.FIELD),
+ INVALID_CATEGORIES);
+ }
+ }
+
+ @Test
+ public void testMixedColumnOrderPreservesMetadataAndValues() throws Exception {
+ final Session session = mock(Session.class);
+ final SessionConnection connection = mock(SessionConnection.class);
+ when(session.getDefaultSessionConnection()).thenReturn(connection);
+ when(session.getDatabase()).thenReturn("db");
+ when(connection.getSessionId()).thenReturn(10L);
+ when(connection.prepareNumericRow(any())).thenReturn(42L);
+ final TableSession tableSession = new TableSession(session);
+ final List<String> names = List.of("v1", "tag1", "v2", "tag2");
+ final List<TSDataType> types =
+ List.of(TSDataType.DOUBLE, TSDataType.STRING, TSDataType.INT32, TSDataType.STRING);
+ final List<ColumnCategory> categories =
+ List.of(ColumnCategory.FIELD, ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.TAG);
+ final NumericRowPreparedStatement prepared =
+ tableSession.prepareInsertNumericRow("t", names, types, categories);
+ final ArgumentCaptor<TSPrepareNumericRowReq> request =
+ ArgumentCaptor.forClass(TSPrepareNumericRowReq.class);
+ verify(connection).prepareNumericRow(request.capture());
+ assertMetadata(prepared, request.getValue());
+ doAnswer(
+ invocation -> {
+ final ByteBuffer values = invocation.getArgument(2);
+ final ByteBuffer tags = invocation.getArgument(3);
+ assertEquals("a", ReadWriteIOUtils.readString(tags));
+ assertEquals("b", ReadWriteIOUtils.readString(tags));
+ assertFalse(tags.hasRemaining());
+ assertEquals(12.5, values.getDouble(), 0.0);
+ assertEquals(7, values.getInt());
+ assertFalse(values.hasRemaining());
+ return null;
+ })
+ .when(connection)
+ .executePreparedNumericRow(eq(42L), eq(123L), any(ByteBuffer.class), any(ByteBuffer.class));
+ tableSession.executePreparedInsertNumericRow(
+ prepared, 123L, new long[] {Double.doubleToRawLongBits(12.5), 7}, new String[] {"a", "b"});
+ verify(connection)
+ .executePreparedNumericRow(eq(42L), eq(123L), any(ByteBuffer.class), any(ByteBuffer.class));
+ assertEquals(List.of("v1", "tag1", "v2", "tag2"), names);
+ assertEquals(
+ List.of(TSDataType.DOUBLE, TSDataType.STRING, TSDataType.INT32, TSDataType.STRING), types);
+ assertEquals(
+ List.of(ColumnCategory.FIELD, ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.TAG),
+ categories);
+ }
+
+ private static void assertMetadata(
+ NumericRowPreparedStatement prepared, TSPrepareNumericRowReq request) {
+ assertEquals("t", prepared.getTableName());
+ assertEquals("db", prepared.getDatabase());
+ assertEquals(2, prepared.getTagCount());
+ assertEquals(List.of("tag1", "tag2", "v1", "v2"), prepared.getColumnNames());
+ assertEquals(
+ List.of(TSDataType.STRING, TSDataType.STRING, TSDataType.DOUBLE, TSDataType.INT32),
+ prepared.getDataTypes());
+ assertEquals(
+ List.of(ColumnCategory.TAG, ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD),
+ prepared.getColumnCategories());
+ assertEquals("t", request.getTableName());
+ assertEquals(prepared.getColumnNames(), request.getMeasurements());
+ assertEquals(
+ prepared.getDataTypes().stream().map(TSDataType::ordinal).toList(), request.getDataTypes());
+ assertEquals(
+ prepared.getColumnCategories().stream().map(c -> (byte) c.ordinal()).toList(),
+ request.getColumnCategories());
+ }
+
+ private static void assertInvalid(
+ String table,
+ List<String> names,
+ List<TSDataType> types,
+ List<ColumnCategory> categories,
+ String message)
+ throws Exception {
+ final Session session = mock(Session.class);
+ try {
+ new TableSession(session).prepareInsertNumericRow(table, names, types, categories);
+ fail("Invalid metadata must be rejected before preparing a statement");
+ } catch (IllegalArgumentException e) {
+ assertEquals(message, e.getMessage());
+ }
+ verifyNoMoreInteractions(session);
+ }
+}
diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/PreparedNumericRowPoolTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/PreparedNumericRowPoolTest.java
new file mode 100644
index 00000000000..f65f761c7b9
--- /dev/null
+++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/PreparedNumericRowPoolTest.java
@@ -0,0 +1,222 @@
+/*
+ * 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.session.pool;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.isession.NumericRow;
+import org.apache.iotdb.isession.NumericRowPreparedStatement;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
+import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedNumericRowReq;
+import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowResp;
+import org.apache.iotdb.session.Session;
+import org.apache.iotdb.session.SessionConnection;
+import org.apache.iotdb.session.i18n.SessionMessages;
+
+import org.apache.thrift.TException;
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.utils.BitMap;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.powermock.reflect.Whitebox;
+
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
+
+public class PreparedNumericRowPoolTest {
+ private Session session;
+ private SessionPool pool;
+ private TableSessionWrapper wrapper;
+ private IClientRPCService.Iface client;
+
+ @Before
+ public void setUp() throws Exception {
+ session = new Session.Builder().enableRedirection(false).enableAutoFetch(false).build();
+ client = mock(IClientRPCService.Iface.class);
+ SessionConnection connection = new SessionConnection("table");
+ Whitebox.setInternalState(connection, "client", client);
+ Whitebox.setInternalState(connection, "sessionId", 11L);
+ Whitebox.setInternalState(session, "defaultSessionConnection", connection);
+ pool = mock(SessionPool.class);
+ pool.sqlDialect = session.getSqlDialect();
+ wrapper = new TableSessionWrapper(session, pool);
+ when(client.prepareNumericRow(any()))
+ .thenReturn(new TSPrepareNumericRowResp(success()).setPreparedId(101))
+ .thenReturn(new TSPrepareNumericRowResp(success()).setPreparedId(102));
+ when(client.executePreparedNumericRow(any())).thenReturn(success());
+ when(client.closePreparedNumericRow(anyLong(), anyLong())).thenReturn(success());
+ }
+
+ @Test
+ public void testClosedWrapperCannotUseReborrowedSession() throws Exception {
+ NumericRowPreparedStatement oldPrepared = prepare(wrapper);
+ wrapper.close();
+ // A new borrower owns the same physical connection; the previous wrapper must stay closed.
+ try (TableSessionWrapper next = new TableSessionWrapper(session, pool);
+ NumericRowPreparedStatement nextPrepared = prepare(next)) {
+ clearInvocations(client);
+ IoTDBConnectionException error =
+ assertThrows(IoTDBConnectionException.class, () -> prepare(wrapper));
+ assertEquals(SessionMessages.EXCEPTION_TABLE_SESSION_IS_CLOSED_77FE359B, error.getMessage());
+ assertThrows(
+ IoTDBConnectionException.class,
+ () ->
+ wrapper.executePreparedInsertNumericRow(
+ oldPrepared, 1L, new long[] {42}, new String[] {"a"}));
+ assertThrows(
+ IoTDBConnectionException.class,
+ () ->
+ wrapper.executePreparedInsertNumericRow(
+ oldPrepared, 1L, new long[] {42}, new String[] {"a"}, new BitMap(1)));
+ verifyZeroInteractions(client);
+ verify(pool, never()).cleanSessionAndMayThrowConnectionException(any());
+
+ next.executePreparedInsertNumericRow(nextPrepared, 2L, new long[] {43}, new String[] {"a"});
+ ArgumentCaptor<TSExecutePreparedNumericRowReq> request =
+ ArgumentCaptor.forClass(TSExecutePreparedNumericRowReq.class);
+ verify(client).executePreparedNumericRow(request.capture());
+ assertEquals(102L, request.getValue().getPreparedId());
+ }
+ wrapper.close();
+ verify(pool, times(2)).putBack(session);
+ }
+
+ @Test
+ public void testClosedWrapperCannotInsertNumericRow() throws Exception {
+ wrapper.close();
+
+ IoTDBConnectionException error =
+ assertThrows(
+ IoTDBConnectionException.class,
+ () ->
+ wrapper.insert(
+ new NumericRow(
+ "t",
+ 1L,
+ List.of("tag", "v"),
+ List.of(TSDataType.STRING, TSDataType.INT64),
+ List.of(ColumnCategory.TAG, ColumnCategory.FIELD),
+ new long[] {42L},
+ new String[] {"a"})));
+
+ assertEquals(SessionMessages.EXCEPTION_TABLE_SESSION_IS_CLOSED_77FE359B, error.getMessage());
+ verifyZeroInteractions(client);
+ }
+
+ @Test
+ public void testPrepareTransportFailureDiscardsSession() throws Exception {
+ TException transportFailure = new TException("transport closed");
+ when(client.prepareNumericRow(any())).thenThrow(transportFailure);
+ IoTDBConnectionException error =
+ assertThrows(
+ IoTDBConnectionException.class,
+ () -> {
+ try (TableSessionWrapper borrowed = wrapper) {
+ prepare(borrowed);
+ }
+ });
+ assertSame(transportFailure, error.getCause());
+ // Automatic close must neither mask the transport failure nor return the broken session.
+ assertEquals(0, error.getSuppressed().length);
+ assertThrows(IoTDBConnectionException.class, () -> prepare(wrapper));
+ wrapper.close();
+ verify(pool).cleanSessionAndMayThrowConnectionException(session);
+ verify(pool, never()).putBack(any());
+ verify(client).prepareNumericRow(any());
+ }
+
+ @Test
+ public void testExecuteTransportFailureDiscardsSession() throws Exception {
+ NumericRowPreparedStatement prepared = prepare(wrapper);
+ TException transportFailure = new TException("transport closed");
+ when(client.executePreparedNumericRow(any())).thenThrow(transportFailure);
+ IoTDBConnectionException error =
+ assertThrows(
+ IoTDBConnectionException.class,
+ () -> {
+ try (TableSessionWrapper borrowed = wrapper) {
+ borrowed.executePreparedInsertNumericRow(
+ prepared, 1L, new long[] {42}, new String[] {"a"});
+ }
+ });
+ assertSame(transportFailure, error.getCause());
+ assertEquals(0, error.getSuppressed().length);
+ assertThrows(
+ IoTDBConnectionException.class,
+ () ->
+ wrapper.executePreparedInsertNumericRow(
+ prepared, 2L, new long[] {43}, new String[] {"a"}));
+ verify(pool).cleanSessionAndMayThrowConnectionException(session);
+ verify(pool, never()).putBack(any());
+ verify(client).executePreparedNumericRow(any());
+ }
+
+ @Test
+ public void testPrepareStatementErrorPreservesHealthySession() throws Exception {
+ when(client.prepareNumericRow(any()))
+ .thenReturn(
+ new TSPrepareNumericRowResp(
+ new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode())
+ .setMessage("invalid prepared metadata")));
+ try (TableSessionWrapper borrowed = wrapper) {
+ assertThrows(StatementExecutionException.class, () -> prepare(borrowed));
+ }
+ verify(pool, never()).cleanSessionAndMayThrowConnectionException(any());
+ verify(pool).putBack(session);
+ }
+
+ @Test
+ public void testCloseReturnsSessionOnlyOnce() throws Exception {
+ wrapper.close();
+ wrapper.close();
+ verify(pool).putBack(session);
+ verify(pool, never()).cleanSessionAndMayThrowConnectionException(any());
+ verifyZeroInteractions(client);
+ }
+
+ private static NumericRowPreparedStatement prepare(TableSessionWrapper borrowed)
+ throws Exception {
+ return borrowed.prepareInsertNumericRow(
+ "t",
+ List.of("tag", "v"),
+ List.of(TSDataType.STRING, TSDataType.INT64),
+ List.of(ColumnCategory.TAG, ColumnCategory.FIELD));
+ }
+
+ private static TSStatus success() {
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatement.java
index 99edef690b5..e83d2afc53b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatement.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatement.java
@@ -26,6 +26,8 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.exception.metadata.PathNotExistException;
import org.apache.iotdb.db.i18n.DataNodePipeMessages;
import org.apache.iotdb.db.pipe.receiver.transform.converter.ValueConverter;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
import org.apache.tsfile.annotations.TableModel;
@@ -45,6 +47,31 @@ public class PipeConvertedInsertRowStatement extends InsertRowStatement {
LoggerFactory.getLogger(PipeConvertedInsertRowStatement.class);
public PipeConvertedInsertRowStatement(final InsertRowStatement insertRowStatement) {
+ this((InsertBaseStatement) insertRowStatement);
+ time = insertRowStatement.getTime();
+ values = insertRowStatement.getValues();
+ isNeedInferType = insertRowStatement.isNeedInferType();
+ deviceID = insertRowStatement.getRawTableDeviceID();
+ removeAllFailedMeasurementMarks();
+ }
+
+ public PipeConvertedInsertRowStatement(final InsertNumericRowStatement insertRowStatement) {
+ this((InsertBaseStatement) insertRowStatement);
+ time = insertRowStatement.getTime();
+ // Box only at the Pipe conversion boundary: the receiver may require non-numeric FIELDs.
+ values = new Object[insertRowStatement.getValueCount()];
+ for (int i = 0; i < values.length; i++) {
+ values[i] = insertRowStatement.getFirstValueOfIndex(i);
+ }
+ isNeedInferType = false;
+ // Schema analysis and failed-column recovery must not mutate the original numeric row.
+ columnCategories = columnCategories == null ? null : columnCategories.clone();
+ tagColumnIndices = null;
+ attrColumnIndices = null;
+ removeAllFailedMeasurementMarks();
+ }
+
+ private PipeConvertedInsertRowStatement(final InsertBaseStatement insertRowStatement) {
super();
// Statement
isDebug = insertRowStatement.isDebug();
@@ -59,12 +86,6 @@ public class PipeConvertedInsertRowStatement extends InsertRowStatement {
attrColumnIndices = insertRowStatement.getAttrColumnIndices();
writeToTable = insertRowStatement.isWriteToTable();
databaseName = insertRowStatement.getDatabaseName().orElse(null);
- // InsertRowStatement
- time = insertRowStatement.getTime();
- values = insertRowStatement.getValues();
- isNeedInferType = insertRowStatement.isNeedInferType();
- deviceID = insertRowStatement.getRawTableDeviceID();
-
// To ensure that the measurement remains unchanged during the WAL writing process, the array
// needs to be copied before the failed Measurement mark can be deleted.
final MeasurementSchema[] measurementSchemas = insertRowStatement.getMeasurementSchemas();
@@ -87,8 +108,6 @@ public class PipeConvertedInsertRowStatement extends InsertRowStatement {
if (failedMeasurementIndex2Info != null) {
this.failedMeasurementIndex2Info = new HashMap<>(failedMeasurementIndex2Info);
}
-
- removeAllFailedMeasurementMarks();
}
@Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/visitor/PipeTableStatementDataTypeConvertExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/visitor/PipeTableStatementDataTypeConvertExecutionVisitor.java
index 5a80da6f6e6..d76428d4813 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/visitor/PipeTableStatementDataTypeConvertExecutionVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/visitor/PipeTableStatementDataTypeConvertExecutionVisitor.java
@@ -34,6 +34,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.StatementNode;
import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsOfOneDeviceStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement;
@@ -221,6 +222,15 @@ public class PipeTableStatementDataTypeConvertExecutionVisitor
new PipeConvertedInsertRowStatement(insertRowStatement), statusAndDatabaseName.getRight());
}
+ @Override
+ public Optional<TSStatus> visitInsertNumericRow(
+ final InsertNumericRowStatement insertNumericRowStatement,
+ final Pair<TSStatus, String> statusAndDatabaseName) {
+ return tryExecute(
+ new PipeConvertedInsertRowStatement(insertNumericRowStatement),
+ statusAndDatabaseName.getRight());
+ }
+
@Override
public Optional<TSStatus> visitInsertRows(
final InsertRowsStatement insertRowsStatement,
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 7435e901d2f..d96df20b784 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
@@ -2614,6 +2614,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
IClientSession session = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
if (!SESSION_MANAGER.checkLogin(session)) return getNotLoggedInStatus();
InsertNumericRowStatement statement = null;
+ OperationQuota quota = null;
try {
if (req.getSessionId() != session.getId()) return RpcUtils.getStatus(TSStatusCode.NOT_LOGIN);
statement =
@@ -2624,6 +2625,8 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
ByteBuffer.wrap(req.getTagValues()),
ByteBuffer.wrap(req.getValues()),
req.bufferForNullBitmap());
+ quota =
+ DataNodeThrottleQuotaManager.getInstance().checkQuota(session.getUsername(), statement);
long queryId = SESSION_MANAGER.requestQueryId();
return COORDINATOR.executeForTableModel(
statement,
@@ -2635,6 +2638,8 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
metadata,
config.getConnectionTimeoutInMS())
.status;
+ } catch (IoTDBException e) {
+ return onIoTDBException(e, OperationType.INSERT_RECORD, e.getErrorCode());
} catch (Exception e) {
return onQueryException(
e, OperationType.INSERT_RECORD.getName(), TSStatusCode.EXECUTE_STATEMENT_ERROR);
@@ -2643,6 +2648,9 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
// Remove only this session's invalid handle, releasing its reservation exactly once.
session.removePreparedNumericRow(req.getPreparedId());
}
+ if (quota != null) {
+ quota.close();
+ }
}
}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatementTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatementTest.java
index 8c563a84b1d..a58a05af70a 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatementTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/transform/statement/PipeConvertedInsertRowStatementTest.java
@@ -19,15 +19,26 @@
package org.apache.iotdb.db.pipe.receiver.transform.statement;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
+import org.apache.iotdb.db.pipe.receiver.visitor.PipeTableStatementDataTypeConvertExecutionVisitor;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
+import org.apache.iotdb.db.utils.NumericRowValues;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.tsfile.common.conf.TSFileConfig;
import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.Pair;
import org.apache.tsfile.write.schema.MeasurementSchema;
import org.junit.Assert;
import org.junit.Test;
import java.time.ZoneId;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
public class PipeConvertedInsertRowStatementTest {
@@ -52,4 +63,139 @@ public class PipeConvertedInsertRowStatementTest {
Assert.assertNull(convertedStatement.getValues()[0]);
Assert.assertFalse(convertedStatement.isNeedInferType());
}
+
+ @Test
+ public void testNumericRowDispatchConvertsNumericAndStringFields() throws Exception {
+ final InsertNumericRowStatement original = numericRow();
+ final long[] originalValues = original.getValues().clone();
+ final AtomicInteger executions = new AtomicInteger();
+ final TSStatus success = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ final PipeTableStatementDataTypeConvertExecutionVisitor visitor =
+ new PipeTableStatementDataTypeConvertExecutionVisitor(
+ (statement, database) -> {
+ executions.incrementAndGet();
+ Assert.assertEquals("db", database);
+ Assert.assertTrue(statement instanceof PipeConvertedInsertRowStatement);
+ final PipeConvertedInsertRowStatement converted =
+ (PipeConvertedInsertRowStatement) statement;
+ Assert.assertTrue(converted.isForceTypeConversion());
+ Assert.assertFalse(converted.isNeedInferType());
+ Assert.assertTrue(converted.isAligned());
+ Assert.assertTrue(converted.isWriteToTable());
+ Assert.assertEquals(original.getTime(), converted.getTime());
+ Assert.assertEquals(original.getDatabaseName(), converted.getDatabaseName());
+ Assert.assertEquals(original.getTableDeviceID(), converted.getTableDeviceID());
+ Assert.assertEquals(original.getFirstValueOfIndex(0), converted.getValues()[0]);
+ Assert.assertTrue(converted.checkAndCastDataType(1, TSDataType.INT32));
+ Assert.assertTrue(converted.checkAndCastDataType(2, TSDataType.STRING));
+ Assert.assertTrue(converted.checkAndCastDataType(3, TSDataType.STRING));
+ Assert.assertEquals(12, converted.getValues()[1]);
+ Assert.assertEquals(
+ new Binary("8.5", TSFileConfig.STRING_CHARSET), converted.getValues()[2]);
+ Assert.assertNull(converted.getValues()[3]);
+ Assert.assertArrayEquals(
+ new TSDataType[] {
+ TSDataType.STRING, TSDataType.INT32, TSDataType.STRING, TSDataType.STRING
+ },
+ converted.getDataTypes());
+ return success;
+ });
+
+ Assert.assertEquals(success, visitor.process(original, new Pair<>(success, "db")).orElse(null));
+ Assert.assertEquals(1, executions.get());
+ Assert.assertArrayEquals(originalValues, original.getValues());
+ Assert.assertArrayEquals(
+ new TSDataType[] {
+ TSDataType.STRING, TSDataType.DOUBLE, TSDataType.DOUBLE, TSDataType.INT64
+ },
+ original.getDataTypes());
+ Assert.assertTrue(original.getNumericValues().isNull(3));
+ }
+
+ @Test
+ public void testNumericRowRetryRestoresFailuresOnlyInCopy() throws Exception {
+ final InsertNumericRowStatement original = numericRow();
+ original.markFailedMeasurement(0, new IllegalArgumentException("tag failure"));
+ original.markFailedMeasurement(1, new IllegalArgumentException("field failure"));
+ final long[] failedValues = original.getValues().clone();
+ final PipeConvertedInsertRowStatement converted = new PipeConvertedInsertRowStatement(original);
+
+ Assert.assertFalse(converted.hasFailedMeasurements());
+ Assert.assertEquals(List.of(0), converted.getTagColumnIndices());
+ Assert.assertEquals("tag", converted.getMeasurements()[0]);
+ Assert.assertEquals("count", converted.getMeasurements()[1]);
+ Assert.assertEquals(new Binary("d1", TSFileConfig.STRING_CHARSET), converted.getValues()[0]);
+ Assert.assertEquals(12.5, converted.getValues()[1]);
+ Assert.assertTrue(converted.checkAndCastDataType(1, TSDataType.INT32));
+ Assert.assertEquals(12, converted.getValues()[1]);
+
+ Assert.assertEquals(2, original.getFailedMeasurementNumber());
+ Assert.assertNotSame(
+ original.getFailedMeasurementInfoMap(), converted.getFailedMeasurementInfoMap());
+ Assert.assertNull(original.getMeasurements()[0]);
+ Assert.assertNull(original.getMeasurements()[1]);
+ Assert.assertNull(original.getDataTypes()[1]);
+ Assert.assertTrue(original.getNumericValues().isNull(0));
+ Assert.assertTrue(original.getNumericValues().isNull(1));
+ Assert.assertArrayEquals(failedValues, original.getValues());
+ Assert.assertEquals(12.5, original.getFailedMeasurementInfoMap().get(1).getValue());
+ }
+
+ @Test
+ public void testNumericRowConversionOwnsMutableMetadataAndPreservesNulls() throws Exception {
+ final InsertNumericRowStatement original = numericRow();
+ original.getNumericValues().setTag(0, null);
+ final PipeConvertedInsertRowStatement converted = new PipeConvertedInsertRowStatement(original);
+ Assert.assertNull(converted.getValues()[0]);
+ Assert.assertNull(converted.getValues()[3]);
+ Assert.assertEquals(original.getTableDeviceID(), converted.getTableDeviceID());
+ Assert.assertNotSame(original.getMeasurements(), converted.getMeasurements());
+ Assert.assertNotSame(original.getDataTypes(), converted.getDataTypes());
+ Assert.assertNotSame(original.getMeasurementSchemas(), converted.getMeasurementSchemas());
+ Assert.assertNotSame(original.getColumnCategories(), converted.getColumnCategories());
+ converted.swapColumn(1, 2);
+ converted.setColumnCategory(TsTableColumnCategory.ATTRIBUTE, 1);
+ converted.markFailedMeasurement(2, new IllegalArgumentException("conversion failure"));
+
+ Assert.assertArrayEquals(
+ new String[] {"tag", "count", "text", "nullable"}, original.getMeasurements());
+ Assert.assertEquals(TsTableColumnCategory.FIELD, original.getColumnCategory(1));
+ Assert.assertEquals("count", original.getMeasurementSchemas()[1].getMeasurementName());
+ Assert.assertEquals(12.5, original.getFirstValueOfIndex(1));
+ Assert.assertEquals(8.5, original.getFirstValueOfIndex(2));
+ Assert.assertFalse(original.hasFailedMeasurements());
+ }
+
+ private static InsertNumericRowStatement numericRow() throws Exception {
+ final InsertNumericRowStatement statement = new InsertNumericRowStatement();
+ statement.setDevicePath(new PartialPath("t"));
+ statement.setDatabaseName("db");
+ statement.setTime(123L);
+ statement.setMeasurements(new String[] {"tag", "count", "text", "nullable"});
+ statement.setDataTypes(
+ new TSDataType[] {
+ TSDataType.STRING, TSDataType.DOUBLE, TSDataType.DOUBLE, TSDataType.INT64
+ });
+ statement.setColumnCategories(
+ new TsTableColumnCategory[] {
+ TsTableColumnCategory.TAG,
+ TsTableColumnCategory.FIELD,
+ TsTableColumnCategory.FIELD,
+ TsTableColumnCategory.FIELD
+ });
+ statement.setMeasurementSchemas(
+ new MeasurementSchema[] {
+ new MeasurementSchema("tag", TSDataType.STRING),
+ new MeasurementSchema("count", TSDataType.DOUBLE),
+ new MeasurementSchema("text", TSDataType.DOUBLE),
+ new MeasurementSchema("nullable", TSDataType.INT64)
+ });
+ final NumericRowValues values = new NumericRowValues(1, 3);
+ values.setTag(0, "d1");
+ values.setObject(1, TSDataType.DOUBLE, 12.5);
+ values.setObject(2, TSDataType.DOUBLE, 8.5);
+ values.setNull(3);
+ statement.setNumericValues(values);
+ return statement;
+ }
}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImplPreparedNumericRowTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImplPreparedNumericRowTest.java
new file mode 100644
index 00000000000..917bcfa4d76
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImplPreparedNumericRowTest.java
@@ -0,0 +1,195 @@
+/*
+ * 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.protocol.thrift.impl;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.exception.RpcThrottlingException;
+import org.apache.iotdb.commons.exception.SemanticException;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+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.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
+import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
+import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement;
+import org.apache.iotdb.db.storageengine.rescon.quotas.DataNodeThrottleQuotaManager;
+import org.apache.iotdb.db.storageengine.rescon.quotas.OperationQuota;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedNumericRowReq;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.InOrder;
+import org.powermock.api.mockito.PowerMockito;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.reflect.Whitebox;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
+
+@RunWith(PowerMockRunner.class)
+@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
+@PrepareForTest({
+ ClientRPCServiceImpl.class,
+ DataNodeThrottleQuotaManager.class,
+ PreparedNumericRowInfo.class
+})
+public class ClientRPCServiceImplPreparedNumericRowTest {
+ private ClientRPCServiceImpl service;
+ private int originalDataNodeId;
+ private Coordinator coordinator;
+ private Coordinator originalCoordinator;
+ private SessionManager originalSessionManager;
+ private IClientSession session;
+ private PreparedNumericRowInfo prepared;
+ private InsertNumericRowStatement statement;
+ private DataNodeThrottleQuotaManager quotaManager;
+ private OperationQuota quota;
+ private TSExecutePreparedNumericRowReq request;
+ private TSStatus success;
+
+ @Before
+ public void setUp() throws Exception {
+ // Coordinator initialization requires a DataNode ID even when query execution is mocked.
+ originalDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
+ IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0);
+ // Exercise the actual RPC method while isolating query execution and per-user quota state.
+ service = mock(ClientRPCServiceImpl.class, CALLS_REAL_METHODS);
+ originalCoordinator = Whitebox.getInternalState(ClientRPCServiceImpl.class, "COORDINATOR");
+ originalSessionManager =
+ Whitebox.getInternalState(ClientRPCServiceImpl.class, "SESSION_MANAGER");
+ coordinator = mock(Coordinator.class);
+ SessionManager sessions = mock(SessionManager.class);
+ Whitebox.setInternalState(ClientRPCServiceImpl.class, "COORDINATOR", coordinator);
+ Whitebox.setInternalState(ClientRPCServiceImpl.class, "SESSION_MANAGER", sessions);
+ session = mock(IClientSession.class);
+ prepared = PowerMockito.mock(PreparedNumericRowInfo.class);
+ statement = new InsertNumericRowStatement();
+ quotaManager = mock(DataNodeThrottleQuotaManager.class);
+ quota = mock(OperationQuota.class);
+ PowerMockito.mockStatic(DataNodeThrottleQuotaManager.class);
+ when(DataNodeThrottleQuotaManager.getInstance()).thenReturn(quotaManager);
+ when(sessions.getCurrSessionAndUpdateIdleTime()).thenReturn(session);
+ when(sessions.checkLogin(session)).thenReturn(true);
+ when(session.getId()).thenReturn(11L);
+ when(session.getUsername()).thenReturn("quota-user");
+ when(session.getPreparedNumericRow(101L)).thenReturn(prepared);
+ when(prepared.bind(anyLong(), any(), any(), any())).thenReturn(statement);
+ when(quotaManager.checkQuota("quota-user", statement)).thenReturn(quota);
+ success = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ when(coordinator.executeForTableModel(
+ same(statement), any(), same(session), anyLong(), any(), eq(""), any(), anyLong()))
+ .thenReturn(new ExecutionResult(new QueryId("prepared_quota"), success));
+ request =
+ new TSExecutePreparedNumericRowReq(
+ 11L, 101L, 1L, ByteBuffer.allocate(0), ByteBuffer.allocate(0));
+ }
+
+ @After
+ public void tearDown() {
+ IoTDBDescriptor.getInstance().getConfig().setDataNodeId(originalDataNodeId);
+ Whitebox.setInternalState(ClientRPCServiceImpl.class, "COORDINATOR", originalCoordinator);
+ Whitebox.setInternalState(
+ ClientRPCServiceImpl.class, "SESSION_MANAGER", originalSessionManager);
+ }
+
+ @Test
+ public void testChecksQuotaBeforeExecutionAndClosesOnSuccess() throws Exception {
+ assertSame(success, service.executePreparedNumericRow(request));
+ InOrder order = inOrder(quotaManager, coordinator, quota);
+ order.verify(quotaManager).checkQuota("quota-user", statement);
+ order
+ .verify(coordinator)
+ .executeForTableModel(
+ same(statement), any(), same(session), anyLong(), any(), eq(""), any(), anyLong());
+ order.verify(quota).close();
+ verify(session, never()).removePreparedNumericRow(anyLong());
+ }
+
+ @Test
+ public void testQuotaRejectionDoesNotExecuteOrInvalidatePreparedRow() throws Exception {
+ int code = TSStatusCode.NUM_WRITE_REQUESTS_EXCEEDED.getStatusCode();
+ when(quotaManager.checkQuota("quota-user", statement))
+ .thenThrow(new RpcThrottlingException("write request quota exceeded", code));
+ TSStatus status = service.executePreparedNumericRow(request);
+ assertEquals(code, status.getCode());
+ // Standard RPC handling adds an operation prefix while retaining the quota failure detail.
+ assertTrue(status.getMessage().endsWith("write request quota exceeded"));
+ verifyZeroInteractions(coordinator, quota);
+ verify(session, never()).removePreparedNumericRow(anyLong());
+ }
+
+ @Test
+ public void testExecutionExceptionClosesQuota() throws Exception {
+ when(coordinator.executeForTableModel(
+ same(statement), any(), same(session), anyLong(), any(), eq(""), any(), anyLong()))
+ .thenThrow(new IllegalStateException("executor failure"));
+ TSStatus status = service.executePreparedNumericRow(request);
+ assertNotEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
+ verify(quota).close();
+ verify(session, never()).removePreparedNumericRow(anyLong());
+ }
+
+ @Test
+ public void testSchemaFailureClosesQuotaAndRemovesPreparedRow() throws Exception {
+ TSStatus schemaFailure = new TSStatus(TSStatusCode.SEMANTIC_ERROR.getStatusCode());
+ when(coordinator.executeForTableModel(
+ same(statement), any(), same(session), anyLong(), any(), eq(""), any(), anyLong()))
+ .thenAnswer(
+ invocation -> {
+ statement.markTableSchemaValidationFailed();
+ return new ExecutionResult(new QueryId("prepared_quota"), schemaFailure);
+ });
+ assertSame(schemaFailure, service.executePreparedNumericRow(request));
+ verify(quotaManager).checkQuota("quota-user", statement);
+ verify(quota).close();
+ verify(session).removePreparedNumericRow(101L);
+ }
+
+ @Test
+ public void testInvalidPayloadDoesNotConsumeQuota() throws Exception {
+ when(prepared.bind(anyLong(), any(), any(), any()))
+ .thenThrow(new SemanticException("invalid numeric payload"));
+ TSStatus status = service.executePreparedNumericRow(request);
+ assertNotEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
+ verifyZeroInteractions(quotaManager, coordinator, quota);
+ verify(session, never()).removePreparedNumericRow(anyLong());
+ }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/NumericTableRowTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/NumericTableRowTest.java
index 70c39d90072..fbb69b4be04 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/NumericTableRowTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/NumericTableRowTest.java
@@ -20,10 +20,18 @@ package org.apache.iotdb.db.queryengine.plan.parser;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.consensus.index.impl.IoTProgressIndex;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
import org.apache.iotdb.commons.queryengine.plan.relational.metadata.ColumnSchema;
+import org.apache.iotdb.commons.request.IConsensusRequest;
import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
+import org.apache.iotdb.consensus.common.request.BatchIndexedConsensusRequest;
+import org.apache.iotdb.consensus.common.request.ByteBufferConsensusRequest;
+import org.apache.iotdb.consensus.common.request.DeserializedBatchIndexedConsensusRequest;
+import org.apache.iotdb.consensus.common.request.IndexedConsensusRequest;
+import org.apache.iotdb.consensus.common.request.IoTConsensusRequest;
+import org.apache.iotdb.db.consensus.statemachine.dataregion.IoTConsensusDataRegionStateMachine;
import org.apache.iotdb.db.pipe.event.common.tablet.parser.TabletInsertionEventTablePatternParser;
import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTabletBatchReqV2;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode;
@@ -35,6 +43,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.StatementType;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
import org.apache.iotdb.db.storageengine.dataregion.memtable.AlignedWritableMemChunkGroup;
import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable;
import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry;
@@ -299,6 +308,58 @@ public class NumericTableRowTest {
}
}
+ @Test
+ public void testIoTConsensusBatchFromMemoryAndWal() throws Exception {
+ // The real state machine must accept both live PlanNode bytes and WAL catch-up envelopes,
+ // attach the sender progress index, and dispatch each numeric row to DataRegion.insert.
+ RelationalInsertNumericRowNode original = nodeWithNullTag();
+ original.markFailedMeasurement(2);
+ original.setSearchIndex(42);
+ original.setLastFragment(true);
+ WALInfoEntry entry = new WALInfoEntry(7, original);
+ ByteBuffer wal = ByteBuffer.allocate(entry.serializedSize());
+ entry.serialize(new WALByteBufferForTest(wal));
+ wal.flip();
+ BatchIndexedConsensusRequest batch = new BatchIndexedConsensusRequest(3);
+ batch.add(
+ new IndexedConsensusRequest(
+ -1,
+ 41,
+ Collections.singletonList(
+ new ByteBufferConsensusRequest(original.serializeToByteBuffer()))));
+ batch.add(
+ new IndexedConsensusRequest(
+ -1, 42, Collections.singletonList(new IoTConsensusRequest(wal))));
+ DataRegion region = Mockito.mock(DataRegion.class);
+ IoTConsensusDataRegionStateMachine machine = new IoTConsensusDataRegionStateMachine(region);
+ DeserializedBatchIndexedConsensusRequest decoded =
+ (DeserializedBatchIndexedConsensusRequest) machine.deserializeRequest(batch);
+ assertEquals(2, decoded.getInsertNodes().size());
+ for (int i = 0; i < 2; i++) {
+ IConsensusRequest request = decoded.getInsertNodes().get(i);
+ assertTrue(request instanceof RelationalInsertNumericRowNode);
+ RelationalInsertNumericRowNode row = (RelationalInsertNumericRowNode) request;
+ assertEquals(original.getDeviceID(), row.getDeviceID());
+ assertEquals(original.getTime(), row.getTime());
+ assertEquals(original.getValidMeasurementNumber(false), row.getMeasurements().length);
+ assertEquals(
+ ((RelationalInsertNumericRowNode) decoded.getInsertNodes().get(0)).getNumericValues(),
+ row.getNumericValues());
+ assertEquals(new IoTProgressIndex(3, 41L + i), row.getProgressIndex());
+ assertObjectStorageAbsent(RelationalInsertNumericRowNode.class, row);
+ }
+ TSStatus status = machine.write(decoded);
+ assertEquals(2, status.getSubStatusSize());
+ for (int i = 0; i < 2; i++) {
+ assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getSubStatus().get(i).getCode());
+ RelationalInsertNumericRowNode row =
+ (RelationalInsertNumericRowNode) decoded.getInsertNodes().get(i);
+ Mockito.verify(region).insert(Mockito.same(row));
+ Mockito.verify(region).insertSeparatorToWAL(Mockito.same(row));
+ }
+ }
+
@Test
public void testPipeBatchRetainsIndependentNumericStatements() throws Exception {
RelationalInsertNumericRowNode node = nodeWithNullTag();