From 02bd5c03d3ecad8ea64a18f76d0e3812d4e2f53d Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Mon, 6 Jul 2026 18:05:23 +0800 Subject: [PATCH 1/9] ver1 --- .../pool/IoTDBInsertTableSessionPoolIT.java | 58 ++++++++++++ .../apache/iotdb/isession/ITableSession.java | 12 +++ .../org/apache/iotdb/session/Session.java | 19 ++++ .../apache/iotdb/session/TableSession.java | 8 ++ .../session/pool/TableSessionWrapper.java | 14 +++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 3 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 3 + .../thrift/impl/ClientRPCServiceImpl.java | 88 +++++++++++++------ .../plan/parser/StatementGenerator.java | 19 ++++ .../analyzer/StatementAnalyzer.java | 6 ++ .../relational/planner/RelationPlanner.java | 27 +++++- .../plan/relational/sql/ast/AstVisitor.java | 4 + .../sql/ast/WrappedInsertStatement.java | 5 +- .../crud/InsertMultiTabletsStatement.java | 7 ++ .../src/main/thrift/client.thrift | 4 +- 15 files changed, 243 insertions(+), 34 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java index c92d09b4dbba6..a5b628cfce838 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java @@ -250,6 +250,64 @@ public void testInsertKeyword() throws IoTDBConnectionException, StatementExecut assertEquals("20240815", rec.getFields().get(12).getStringValue()); } assertFalse(rs1.hasNext()); + + Tablet tablet2 = + new Tablet( + "table20", + IMeasurementSchema.getMeasurementNameList(schemas), + schemas.stream().map(IMeasurementSchema::getType).collect(Collectors.toList()), + columnTypes, + 10); + + for (long row = 10; row < 20; row++) { + int rowIndex = tablet2.getRowSize(); + tablet2.addTimestamp(rowIndex, timestamp + row); + tablet2.addValue("device_id", rowIndex, "2"); + tablet2.addValue("attribute", rowIndex, "2"); + tablet2.addValue("boolean", rowIndex, true); + tablet2.addValue("int32", rowIndex, Integer.valueOf("2")); + tablet2.addValue("int64", rowIndex, Long.valueOf("2")); + tablet2.addValue("float", rowIndex, Float.valueOf("2.0")); + tablet2.addValue("double", rowIndex, Double.valueOf("2.0")); + tablet2.addValue("text", rowIndex, "false"); + tablet2.addValue("string", rowIndex, "false"); + tablet2.addValue("blob", rowIndex, new Binary("iotdb", Charset.defaultCharset())); + tablet2.addValue("timestamp", rowIndex, 2L); + tablet2.addValue("date", rowIndex, LocalDate.parse("2024-08-16")); + } + + Tablet tablet3 = + new Tablet( + "table20", + IMeasurementSchema.getMeasurementNameList(schemas), + schemas.stream().map(IMeasurementSchema::getType).collect(Collectors.toList()), + columnTypes, + 10); + + for (long row = 20; row < 30; row++) { + int rowIndex = tablet3.getRowSize(); + tablet3.addTimestamp(rowIndex, timestamp + row); + tablet3.addValue("device_id", rowIndex, "2"); + tablet3.addValue("attribute", rowIndex, "2"); + tablet3.addValue("boolean", rowIndex, true); + tablet3.addValue("int32", rowIndex, Integer.valueOf("2")); + tablet3.addValue("int64", rowIndex, Long.valueOf("2")); + tablet3.addValue("float", rowIndex, Float.valueOf("2.0")); + tablet3.addValue("double", rowIndex, Double.valueOf("2.0")); + tablet3.addValue("text", rowIndex, "false"); + tablet3.addValue("string", rowIndex, "false"); + tablet3.addValue("blob", rowIndex, new Binary("iotdb", Charset.defaultCharset())); + tablet3.addValue("timestamp", rowIndex, 2L); + tablet3.addValue("date", rowIndex, LocalDate.parse("2024-08-16")); + } + session.insertTablets(Arrays.asList(tablet2, tablet3)); + + try (SessionDataSet rs2 = + session.executeQueryStatement("select count(*) from table20 where device_id='2'")) { + RowRecord rec = rs2.next(); + assertEquals(20, rec.getFields().get(0).getLongV()); + assertFalse(rs2.hasNext()); + } } } } diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java index 75a37f8384719..bed22ed771816 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java @@ -24,6 +24,8 @@ import org.apache.tsfile.write.record.Tablet; +import java.util.List; + /** * This interface defines a session for interacting with IoTDB tables. It supports operations such * as data insertion, executing queries, and closing the session. Implementations of this interface @@ -46,6 +48,16 @@ public interface ITableSession extends AutoCloseable { */ void insert(Tablet tablet) throws StatementExecutionException, IoTDBConnectionException; + /** + * Inserts multiple {@link Tablet}s into the database. + * + * @param tablets the tablets containing time-series data to be inserted. + * @throws StatementExecutionException if an error occurs while executing the statement. + * @throws IoTDBConnectionException if there is an issue with the IoTDB connection. + */ + void insertTablets(List tablets) + throws StatementExecutionException, IoTDBConnectionException; + /** * Executes a non-query SQL statement, such as a DDL or DML command. * diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 6819325a9b27b..4d91199f7df1b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -2857,6 +2857,25 @@ public void insertRelationalTablet(Tablet tablet) } } + /** + * insert relational Tablets. Note: This method is for internal use only, we do not guarantee + * compatibility with subsequent versions. + * + * @param tablets data batches + */ + public void insertRelationalTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException { + TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false); + request.setWriteToTable(true); + for (Tablet tablet : tablets) { + request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes())); + } + try { + getDefaultSessionConnection().insertTablets(request); + } catch (RedirectException ignored) { + } + } + private void insertRelationalTabletWithLeaderCache(Tablet tablet) throws IoTDBConnectionException, StatementExecutionException { Map relationalTabletGroup = new HashMap<>(); 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 059fdea346948..b4730c78c0a3d 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 @@ -26,6 +26,8 @@ import org.apache.tsfile.write.record.Tablet; +import java.util.List; + public class TableSession implements ITableSession { private final Session session; @@ -39,6 +41,12 @@ public void insert(Tablet tablet) throws StatementExecutionException, IoTDBConne session.insertRelationalTablet(tablet); } + @Override + public void insertTablets(List tablets) + throws StatementExecutionException, IoTDBConnectionException { + session.insertRelationalTablets(tablets); + } + @Override public void executeNonQueryStatement(String sql) throws IoTDBConnectionException, StatementExecutionException { 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 4cf7ca980f347..f2f89dd07bbf0 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 @@ -30,6 +30,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.List; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; @@ -68,6 +69,19 @@ public void insert(Tablet tablet) throws StatementExecutionException, IoTDBConne } } + @Override + public void insertTablets(List tablets) + throws StatementExecutionException, IoTDBConnectionException { + try { + session.insertRelationalTablets(tablets); + } catch (IoTDBConnectionException e) { + sessionPool.cleanSessionAndMayThrowConnectionException(session); + closed.set(true); + session = null; + throw e; + } + } + @Override public SessionDataSet executeQueryStatement(String sql) throws StatementExecutionException, IoTDBConnectionException { diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 061b049701191..8d66348a9d5bf 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -1816,6 +1816,9 @@ public final class DataNodeQueryMessages { "' must be consistent across the entire cluster and only one can be set at a time."; public static final String CANNOT_INSERT_INTO_MULTIPLE_DATABASES_WITHIN_ONE_STATEMENT_PLEASE_SPLIT_THEM_MANUALLY = "Cannot insert into multiple databases within one statement, please split them manually"; + public static final String + CANNOT_INSERT_INTO_MULTIPLE_TABLES_WITHIN_ONE_STATEMENT_PLEASE_SPLIT_THEM_MANUALLY_5967C460 = + "Cannot insert into multiple tables within one statement, please split them manually"; public static final String THE_MEASUREMENTLIST_S_SIZE_D_IS_NOT_CONSISTENT_WITH_THE_VALUELIST_S_SIZE_D = "the measurementList's size %d is not consistent with the valueList's size %d"; public static final String THE_MEASUREMENTLIST_S_SIZE_D_IS_NOT_CONSISTENT_WITH_THE_COLUMNLIST_S_SIZE_D = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 1d7df89e0278b..7356b65bfb611 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -1960,6 +1960,9 @@ public final class DataNodeQueryMessages { "' 必须在整个集群中保持一致，且同一时间只能设置一个。"; public static final String CANNOT_INSERT_INTO_MULTIPLE_DATABASES_WITHIN_ONE_STATEMENT_PLEASE_SPLIT_THEM_MANUALLY = "不能在一条语句中写入多个数据库，请手动拆分。"; + public static final String + CANNOT_INSERT_INTO_MULTIPLE_TABLES_WITHIN_ONE_STATEMENT_PLEASE_SPLIT_THEM_MANUALLY_5967C460 = + "不能在一条语句中写入多个表，请手动拆分。"; public static final String THE_MEASUREMENTLIST_S_SIZE_D_IS_NOT_CONSISTENT_WITH_THE_VALUELIST_S_SIZE_D = "measurementList 的大小 %d 与 valueList 的大小 %d 不一致"; public static final String THE_MEASUREMENTLIST_S_SIZE_D_IS_NOT_CONSISTENT_WITH_THE_COLUMNLIST_S_SIZE_D = 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 b6541a4f0a0dc..6bdb10c3fc25c 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 @@ -2416,14 +2416,28 @@ public TSStatus insertRecords(TSInsertRecordsReq req) { // Step 2: call the coordinator long queryId = SESSION_MANAGER.requestQueryId(); - ExecutionResult result = - COORDINATOR.executeForTreeModel( - statement, - queryId, - SESSION_MANAGER.getSessionInfo(clientSession), - "", - partitionFetcher, - schemaFetcher); + ExecutionResult result; + if (statement.isWriteToTable()) { + result = + COORDINATOR.executeForTableModel( + statement, + relationSqlParser, + clientSession, + queryId, + SESSION_MANAGER.getSessionInfo(clientSession), + "", + metadata, + config.getConnectionTimeoutInMS()); + } else { + result = + COORDINATOR.executeForTreeModel( + statement, + queryId, + SESSION_MANAGER.getSessionInfo(clientSession), + "", + partitionFetcher, + schemaFetcher); + } return result.status; } catch (IoTDBException e) { @@ -2667,8 +2681,10 @@ public TSStatus insertTablets(TSInsertTabletsReq req) { if (!SESSION_MANAGER.checkLogin(clientSession)) { return getNotLoggedInStatus(); } - req.setMeasurementsList( - PathUtils.checkIsLegalSingleMeasurementListsAndUpdate(req.getMeasurementsList())); + if (!req.isWriteToTable()) { + req.setMeasurementsList( + PathUtils.checkIsLegalSingleMeasurementListsAndUpdate(req.getMeasurementsList())); + } // Step 1: transfer from TSInsertTabletsReq to Statement InsertMultiTabletsStatement statement = StatementGenerator.createStatement(req); @@ -2677,16 +2693,18 @@ public TSStatus insertTablets(TSInsertTabletsReq req) { return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); } - // permission check - TSStatus status = - AuthorityChecker.checkAuthority( - statement, - new TreeAccessCheckContext( - clientSession.getUserId(), - clientSession.getUsername(), - clientSession.getClientAddress())); - if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { - return status; + // permission check for tree model, table model does this in the analysis stage + if (!req.isWriteToTable()) { + TSStatus status = + AuthorityChecker.checkAuthority( + statement, + new TreeAccessCheckContext( + clientSession.getUserId(), + clientSession.getUsername(), + clientSession.getClientAddress())); + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + return status; + } } quota = @@ -2695,14 +2713,28 @@ public TSStatus insertTablets(TSInsertTabletsReq req) { // Step 2: call the coordinator long queryId = SESSION_MANAGER.requestQueryId(); - ExecutionResult result = - COORDINATOR.executeForTreeModel( - statement, - queryId, - SESSION_MANAGER.getSessionInfo(clientSession), - "", - partitionFetcher, - schemaFetcher); + ExecutionResult result; + if (statement.isWriteToTable()) { + result = + COORDINATOR.executeForTableModel( + statement, + relationSqlParser, + clientSession, + queryId, + SESSION_MANAGER.getSessionInfo(clientSession), + "", + metadata, + config.getConnectionTimeoutInMS()); + } else { + result = + COORDINATOR.executeForTreeModel( + statement, + queryId, + SESSION_MANAGER.getSessionInfo(clientSession), + "", + partitionFetcher, + schemaFetcher); + } return result.status; } catch (IoTDBException e) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java index 6bdc4814a6da6..4af72cafefe2a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGenerator.java @@ -425,12 +425,31 @@ public static InsertMultiTabletsStatement createStatement(TSInsertTabletsReq req } insertTabletStatement.setDataTypes(dataTypes); insertTabletStatement.setAligned(req.isAligned); + insertTabletStatement.setWriteToTable(req.isWriteToTable()); + if (req.isWriteToTable()) { + if (!req.isSetColumnCategoriesList() + || req.getColumnCategoriesListSize() != req.getMeasurementsListSize() + || req.getColumnCategoriesList().get(i).size() != req.measurementsList.get(i).size()) { + throw new IllegalArgumentException( + DataNodeQueryMessages + .QUERY_EXCEPTION_MISSING_OR_INVALID_COLUMN_CATEGORIES_FOR_TABLE_INSERTION_5DF990B9); + } + TsTableColumnCategory[] columnCategories = + new TsTableColumnCategory[req.getColumnCategoriesList().get(i).size()]; + for (int j = 0; j < columnCategories.length; j++) { + columnCategories[j] = + TsTableColumnCategory.fromTsFileColumnCategory( + ColumnCategory.values()[req.getColumnCategoriesList().get(i).get(j).intValue()]); + } + insertTabletStatement.setColumnCategories(columnCategories); + } // skip empty tablet if (insertTabletStatement.isEmpty()) { continue; } insertTabletStatementList.add(insertTabletStatement); } + insertStatement.setWriteToTable(req.isWriteToTable()); insertStatement.setInsertTabletStatementList(insertTabletStatementList); PERFORMANCE_OVERVIEW_METRICS.recordParseCost(System.nanoTime() - startTime); return insertStatement; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java index d20e7a42b0a85..22f6037ba8a9b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java @@ -183,6 +183,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRow; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRows; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablet; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablets; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.PipeEnriched; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Property; @@ -1226,6 +1227,11 @@ public Scope visitInsertTablet(InsertTablet insert, Optional scope) { return visitInsert(insert, scope); } + @Override + public Scope visitInsertTablets(InsertTablets insert, Optional scope) { + return visitInsert(insert, scope); + } + @Override public Scope visitInsertRows(InsertRows node, Optional context) { return visitInsert(node, context); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java index f64f3d20b8fc7..7d87d74b1b612 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java @@ -97,6 +97,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedDeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedInsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedWritePlanNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertMultiTabletsNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalDeleteDataNode; @@ -133,8 +134,10 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRow; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRows; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablet; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablets; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.PipeEnriched; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; @@ -1344,12 +1347,18 @@ public RelationPlan visitValues(Values node, Void context) { @Override public RelationPlan visitInsertTablet(InsertTablet node, Void context) { final InsertTabletStatement insertTabletStatement = node.getInnerTreeStatement(); + RelationalInsertTabletNode insertNode = fromInsertTabletStatement(insertTabletStatement); + return new RelationPlan( + insertNode, analysis.getRootScope(), Collections.emptyList(), outerContext); + } + public RelationalInsertTabletNode fromInsertTabletStatement( + final InsertTabletStatement insertTabletStatement) { String[] measurements = insertTabletStatement.getMeasurements(); MeasurementSchema[] measurementSchemas = insertTabletStatement.getMeasurementSchemas(); stayConsistent(measurements, measurementSchemas); - RelationalInsertTabletNode insertNode = + final RelationalInsertTabletNode insertNode = new RelationalInsertTabletNode( idAllocator.genPlanNodeId(), insertTabletStatement.getDevicePath(), @@ -1366,8 +1375,22 @@ public RelationPlan visitInsertTablet(InsertTablet node, Void context) { if (insertTabletStatement.isSingleDevice()) { insertNode.setSingleDevice(); } + return insertNode; + } + + @Override + public RelationPlan visitInsertTablets(InsertTablets node, Void context) { + final InsertMultiTabletsStatement insertMultiTabletsStatement = node.getInnerTreeStatement(); + final InsertMultiTabletsNode insertMultiTabletsNode = + new InsertMultiTabletsNode(idAllocator.genPlanNodeId()); + for (int i = 0; i < insertMultiTabletsStatement.getInsertTabletStatementList().size(); i++) { + insertMultiTabletsNode.addInsertTabletNode( + fromInsertTabletStatement( + insertMultiTabletsStatement.getInsertTabletStatementList().get(i)), + i); + } return new RelationPlan( - insertNode, analysis.getRootScope(), Collections.emptyList(), outerContext); + insertMultiTabletsNode, analysis.getRootScope(), Collections.emptyList(), outerContext); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AstVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AstVisitor.java index 86084ff2c6a1d..64f234897c781 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AstVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AstVisitor.java @@ -203,6 +203,10 @@ default R visitInsertTablet(InsertTablet node, C context) { return visitStatement(node, context); } + default R visitInsertTablets(InsertTablets node, C context) { + return visitStatement(node, context); + } + default R visitFlush(Flush node, C context) { return visitStatement(node, context); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/WrappedInsertStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/WrappedInsertStatement.java index a56a8f4b6ca2c..ad67db1831dbc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/WrappedInsertStatement.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/WrappedInsertStatement.java @@ -35,7 +35,6 @@ import org.apache.iotdb.db.queryengine.plan.relational.metadata.ITableDeviceSchemaValidation; import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement; -import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; import org.apache.iotdb.db.utils.TypeInferenceUtils; import org.apache.iotdb.rpc.TSStatusCode; @@ -154,7 +153,7 @@ public void validateTableSchema(Metadata metadata, MPPQueryContext context) { public void validateTableSchema( final Metadata metadata, final MPPQueryContext context, - final InsertRowStatement insertRowStatement, + final InsertBaseStatement insertRowStatement, final String databaseName, final boolean allowCreateTable) { metadata.validateInsertNodeMeasurements( @@ -179,7 +178,7 @@ private void validate( validate(index, innerTreeStatement, measurement, dataType, columnCategory, existingColumn); } - private void validate( + void validate( final int index, final InsertBaseStatement innerTreeStatement, final String measurement, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertMultiTabletsStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertMultiTabletsStatement.java index 63e942312d0b6..82acfee1d6ca2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertMultiTabletsStatement.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertMultiTabletsStatement.java @@ -25,6 +25,7 @@ import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaValidation; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablets; import org.apache.iotdb.db.queryengine.plan.statement.StatementType; import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; import org.apache.iotdb.db.schemaengine.schemaregion.attribute.update.UpdateDetailContainer; @@ -158,6 +159,12 @@ public InsertBaseStatement removeLogicalView() { return splitResult; } + @Override + public org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement + toRelationalStatement(MPPQueryContext context) { + return new InsertTablets(this, context); + } + @TableModel @Override public void toLowerCase() { diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift index 7ff12ceb02154..7ee3ede6ff06e 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/client.thrift @@ -285,6 +285,8 @@ struct TSInsertTabletsReq { 6: required list> typesList 7: required list sizeList 8: optional bool isAligned + 9: optional bool writeToTable + 10: optional list> columnCategoriesList } struct TSInsertRecordsReq { @@ -705,4 +707,4 @@ service IClientRPCService { /** For other node's call */ common.TSStatus testConnectionEmptyRPC(); -} \ No newline at end of file +} From 521f9d7b6d7605934d764d3df5fd24c97f287962 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Mon, 6 Jul 2026 18:33:27 +0800 Subject: [PATCH 2/9] support schema validation --- .../plan/analyze/schema/SchemaValidator.java | 26 ++++++ .../analyzer/StatementAnalyzer.java | 15 +++- .../relational/sql/ast/InsertTablets.java | 90 +++++++++++++++++++ 3 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/SchemaValidator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/SchemaValidator.java index ff6efc8b3c1fb..f400ec4f3cd01 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/SchemaValidator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/SchemaValidator.java @@ -27,11 +27,13 @@ import org.apache.iotdb.db.queryengine.common.schematree.ISchemaTree; import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; import org.apache.iotdb.db.queryengine.plan.relational.security.AccessControl; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertTablets; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.WrappedInsertStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsOfOneDeviceStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.CompressionType; @@ -85,6 +87,30 @@ public static void validate( } } + public static void validate( + final Metadata metadata, + final InsertTablets insertStatement, + final MPPQueryContext context, + final AccessControl accessControl) { + try { + for (InsertTabletStatement tabletStatement : + insertStatement.getInnerTreeStatement().getInsertTabletStatementList()) { + accessControl.checkCanInsertIntoTable( + context.getSession().getUserName(), + new QualifiedObjectName( + unQualifyDatabaseName(insertStatement.getDatabase()), + tabletStatement.getDevicePath().getFullPath()), + context); + } + insertStatement.validateTableSchema(metadata, context); + insertStatement.updateAfterSchemaValidation(context); + insertStatement.validateDeviceSchema(metadata, context); + insertStatement.removeAttributeColumns(); + } catch (final QueryProcessException e) { + throw new SemanticException(e.getMessage()); + } + } + public static ISchemaTree validate( ISchemaFetcher schemaFetcher, List devicePaths, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java index 22f6037ba8a9b..ae13960f230ff 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java @@ -1229,7 +1229,10 @@ public Scope visitInsertTablet(InsertTablet insert, Optional scope) { @Override public Scope visitInsertTablets(InsertTablets insert, Optional scope) { - return visitInsert(insert, scope); + return visitInsert( + insert, + scope, + () -> SchemaValidator.validate(metadata, insert, insert.getContext(), accessControl)); } @Override @@ -1238,6 +1241,14 @@ public Scope visitInsertRows(InsertRows node, Optional context) { } private Scope visitInsert(WrappedInsertStatement insert, Optional scope) { + return visitInsert( + insert, + scope, + () -> SchemaValidator.validate(metadata, insert, insert.getContext(), accessControl)); + } + + private Scope visitInsert( + WrappedInsertStatement insert, Optional scope, Runnable schemaValidation) { final Scope ret = Scope.create(); final MPPQueryContext context = insert.getContext(); @@ -1248,7 +1259,7 @@ private Scope visitInsert(WrappedInsertStatement insert, Optional scope) AnalyzeUtils.analyzeInsert( context, innerInsert, - () -> SchemaValidator.validate(metadata, insert, context, accessControl), + schemaValidation, metadata::getOrCreateDataPartition, AnalyzeUtils::computeTableDataPartitionParams, analysis, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java new file mode 100644 index 0000000000000..fbccd7800d028 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.relational.sql.ast; + +import org.apache.iotdb.calc.exception.QueryProcessException; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IAstVisitor; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.plan.analyze.AnalyzeUtils; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; + +import java.util.List; + +public class InsertTablets extends WrappedInsertStatement { + + public InsertTablets( + InsertMultiTabletsStatement insertMultiTabletsStatement, MPPQueryContext context) { + super(insertMultiTabletsStatement, context); + } + + @Override + public R accept(IAstVisitor visitor, C context) { + return ((AstVisitor) visitor).visitInsertTablets(this, context); + } + + @Override + public InsertMultiTabletsStatement getInnerTreeStatement() { + return ((InsertMultiTabletsStatement) super.getInnerTreeStatement()); + } + + @Override + public void updateAfterSchemaValidation(MPPQueryContext context) throws QueryProcessException { + getInnerTreeStatement().updateAfterSchemaValidation(context); + } + + @Override + public String getTableName() { + throw new UnsupportedOperationException(); + } + + @Override + public List getDeviceIdList() { + throw new UnsupportedOperationException(); + } + + @Override + public List getAttributeColumnNameList() { + throw new UnsupportedOperationException(); + } + + @Override + public List getAttributeValueList() { + throw new UnsupportedOperationException(); + } + + @Override + public void validateTableSchema(Metadata metadata, MPPQueryContext context) { + for (InsertTabletStatement insertTabletStatement : + getInnerTreeStatement().getInsertTabletStatementList()) { + final String database = AnalyzeUtils.getDatabaseName(insertTabletStatement, context); + super.validateTableSchema(metadata, context, insertTabletStatement, database, false); + } + } + + @Override + public void validateDeviceSchema(Metadata metadata, MPPQueryContext context) { + for (InsertTabletStatement insertTabletStatement : + getInnerTreeStatement().getInsertTabletStatementList()) { + metadata.validateDeviceSchema(new InsertTablet(insertTabletStatement, context), context); + } + } +} From 67d01c412e8808695b9d3d23e334f804510f0aba Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Mon, 6 Jul 2026 18:45:33 +0800 Subject: [PATCH 3/9] add IT --- .../pool/IoTDBInsertTableSessionPoolIT.java | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java index a5b628cfce838..c3a1cd34b1280 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java @@ -310,4 +310,127 @@ public void testInsertKeyword() throws IoTDBConnectionException, StatementExecut } } } + + @Test + public void testInsertTabletsWithDifferentTables() + throws IoTDBConnectionException, StatementExecutionException { + final ITableSessionPool sessionPool = EnvFactory.getEnv().getTableSessionPool(1); + try (final ITableSession session = sessionPool.getSession()) { + session.executeNonQueryStatement("USE \"test\""); + session.executeNonQueryStatement("DROP TABLE IF EXISTS multi_tablet_table1"); + session.executeNonQueryStatement("DROP TABLE IF EXISTS multi_tablet_table2"); + session.executeNonQueryStatement("DROP TABLE IF EXISTS multi_tablet_table3"); + session.executeNonQueryStatement( + "CREATE TABLE multi_tablet_table1 (tag1 STRING TAG, attr1 STRING ATTRIBUTE, s1 INT64 FIELD)"); + session.executeNonQueryStatement( + "CREATE TABLE multi_tablet_table2 (tag2 STRING TAG, s2 DOUBLE FIELD)"); + session.executeNonQueryStatement( + "CREATE TABLE multi_tablet_table3 (region STRING TAG, plant STRING TAG, status STRING ATTRIBUTE, s3 BOOLEAN FIELD)"); + + final Tablet tablet1 = createTabletForTable1(0, 3); + final Tablet tablet2 = createTabletForTable2(10, 3); + final Tablet tablet3 = createTabletForTable1(3, 2); + final Tablet tablet4 = createTabletForTable3(20, 4); + + session.insertTablets(Arrays.asList(tablet1, tablet2, tablet3, tablet4)); + + assertCount(session, "SELECT COUNT(*) FROM multi_tablet_table1 WHERE tag1 = 'tag1'", 5); + assertCount(session, "SELECT COUNT(*) FROM multi_tablet_table2 WHERE tag2 = 'tag2'", 3); + assertCount( + session, + "SELECT COUNT(*) FROM multi_tablet_table3 WHERE region = 'r1' AND plant = 'p1'", + 4); + assertCount( + session, + "SELECT COUNT(*) FROM multi_tablet_table1 WHERE tag1 = 'tag1' AND attr1 = 'attr1'", + 5); + assertCount( + session, + "SELECT COUNT(*) FROM multi_tablet_table3 WHERE region = 'r1' AND plant = 'p1' AND status = 'running'", + 4); + } + } + + private static Tablet createTabletForTable1(final long startTime, final int rowCount) { + final List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema("tag1", TSDataType.STRING)); + schemas.add(new MeasurementSchema("attr1", TSDataType.STRING)); + schemas.add(new MeasurementSchema("s1", TSDataType.INT64)); + final List columnTypes = + Arrays.asList(ColumnCategory.TAG, ColumnCategory.ATTRIBUTE, ColumnCategory.FIELD); + final Tablet tablet = + new Tablet( + "multi_tablet_table1", + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + columnTypes, + rowCount); + for (int i = 0; i < rowCount; i++) { + final int rowIndex = tablet.getRowSize(); + tablet.addTimestamp(rowIndex, startTime + i); + tablet.addValue("tag1", rowIndex, "tag1"); + tablet.addValue("attr1", rowIndex, "attr1"); + tablet.addValue("s1", rowIndex, startTime + i); + } + return tablet; + } + + private static Tablet createTabletForTable2(final long startTime, final int rowCount) { + final List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema("tag2", TSDataType.STRING)); + schemas.add(new MeasurementSchema("s2", TSDataType.DOUBLE)); + final List columnTypes = + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD); + final Tablet tablet = + new Tablet( + "multi_tablet_table2", + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + columnTypes, + rowCount); + for (int i = 0; i < rowCount; i++) { + final int rowIndex = tablet.getRowSize(); + tablet.addTimestamp(rowIndex, startTime + i); + tablet.addValue("tag2", rowIndex, "tag2"); + tablet.addValue("s2", rowIndex, (startTime + i) * 1.0); + } + return tablet; + } + + private static Tablet createTabletForTable3(final long startTime, final int rowCount) { + final List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema("region", TSDataType.STRING)); + schemas.add(new MeasurementSchema("plant", TSDataType.STRING)); + schemas.add(new MeasurementSchema("status", TSDataType.STRING)); + schemas.add(new MeasurementSchema("s3", TSDataType.BOOLEAN)); + final List columnTypes = + Arrays.asList( + ColumnCategory.TAG, ColumnCategory.TAG, ColumnCategory.ATTRIBUTE, ColumnCategory.FIELD); + final Tablet tablet = + new Tablet( + "multi_tablet_table3", + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + columnTypes, + rowCount); + for (int i = 0; i < rowCount; i++) { + final int rowIndex = tablet.getRowSize(); + tablet.addTimestamp(rowIndex, startTime + i); + tablet.addValue("region", rowIndex, "r1"); + tablet.addValue("plant", rowIndex, "p1"); + tablet.addValue("status", rowIndex, "running"); + tablet.addValue("s3", rowIndex, true); + } + return tablet; + } + + private static void assertCount( + final ITableSession session, final String query, final long expectedCount) + throws IoTDBConnectionException, StatementExecutionException { + try (final SessionDataSet dataSet = session.executeQueryStatement(query)) { + final RowRecord record = dataSet.next(); + assertEquals(expectedCount, record.getFields().get(0).getLongV()); + assertFalse(dataSet.hasNext()); + } + } } From 124f2c3a24ce2d8580a16902627b16a37d77a01f Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 09:59:01 +0800 Subject: [PATCH 4/9] add switch of merge tablets --- .../apache/iotdb/isession/SessionConfig.java | 1 + .../iotdb/session/AbstractSessionBuilder.java | 1 + .../org/apache/iotdb/session/Session.java | 224 +++++++++++++++++- .../iotdb/session/TableSessionBuilder.java | 12 + .../iotdb/session/pool/SessionPool.java | 9 + .../session/pool/TableSessionPoolBuilder.java | 12 + .../org/apache/iotdb/session/SessionTest.java | 223 +++++++++++++++++ 7 files changed, 480 insertions(+), 2 deletions(-) diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionConfig.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionConfig.java index 7885d50cf5da1..5fcfd4d9892a6 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionConfig.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/SessionConfig.java @@ -31,6 +31,7 @@ public class SessionConfig { public static final int DEFAULT_CONNECTION_TIMEOUT_MS = 0; public static final boolean DEFAULT_REDIRECTION_MODE = true; public static final boolean DEFAULT_RECORDS_AUTO_CONVERT_TABLET = true; + public static final boolean DEFAULT_ENABLE_MERGE_TABLETS = true; public static final int CPU_CORES = Runtime.getRuntime().availableProcessors(); public static final int DEFAULT_SESSION_EXECUTOR_THREAD_NUM = 2 * CPU_CORES; public static final int DEFAULT_SESSION_EXECUTOR_TASK_NUM = 1_000; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java index 44500ed3f10a6..5b557ae68035d 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java @@ -53,6 +53,7 @@ public abstract class AbstractSessionBuilder { // datanode while encountering retriable errors in current DataNode public boolean enableRedirection = SessionConfig.DEFAULT_REDIRECTION_MODE; public boolean enableRecordsAutoConvertTablet = SessionConfig.DEFAULT_RECORDS_AUTO_CONVERT_TABLET; + public boolean enableMergeTablets = SessionConfig.DEFAULT_ENABLE_MERGE_TABLETS; public Version version = SessionConfig.DEFAULT_VERSION; public long timeOut = SessionConfig.DEFAULT_QUERY_TIME_OUT; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 4d91199f7df1b..8e66fa7f8e573 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -65,6 +65,7 @@ import org.apache.thrift.TException; import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.enums.CompressionType; @@ -166,9 +167,12 @@ public class Session implements ISession { protected boolean enableRedirection; protected boolean enableRecordsAutoConvertTablet = SessionConfig.DEFAULT_RECORDS_AUTO_CONVERT_TABLET; + protected boolean enableMergeTablets = SessionConfig.DEFAULT_ENABLE_MERGE_TABLETS; private static final double CONVERT_THRESHOLD = 0.5; private static final double SAMPLE_PROPORTION = 0.05; private static final int MIN_RECORDS_SIZE = 40; + private static final int MAX_CONSECUTIVE_RELATIONAL_TABLET_MERGE_MISS_COUNT = 10; + private int consecutiveRelationalTabletMergeMissCount = 0; @SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile" protected volatile Map deviceIdToEndpoint; @@ -465,6 +469,7 @@ public Session(AbstractSessionBuilder builder) { this.columnEncodersMap = builder.columnEncodersMap; this.enableRedirection = builder.enableRedirection; this.enableRecordsAutoConvertTablet = builder.enableRecordsAutoConvertTablet; + this.enableMergeTablets = builder.enableMergeTablets; this.username = builder.username; this.password = builder.pw; this.useEncryptedPassword = builder.useEncryptedPassword; @@ -2858,13 +2863,15 @@ public void insertRelationalTablet(Tablet tablet) } /** - * insert relational Tablets. Note: This method is for internal use only, we do not guarantee - * compatibility with subsequent versions. + * insert relational Tablets. * * @param tablets data batches */ public void insertRelationalTablets(List tablets) throws IoTDBConnectionException, StatementExecutionException { + if (enableMergeTablets) { + tablets = mergeRelationalTablets(tablets); + } TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false); request.setWriteToTable(true); for (Tablet tablet : tablets) { @@ -2876,6 +2883,214 @@ public void insertRelationalTablets(List tablets) } } + private List mergeRelationalTablets(final List tablets) { + if (consecutiveRelationalTabletMergeMissCount + >= MAX_CONSECUTIVE_RELATIONAL_TABLET_MERGE_MISS_COUNT) { + return tablets; + } + final List sortedTablets = new ArrayList<>(tablets); + sortedTablets.sort( + Comparator.comparing( + Tablet::getTableName, Comparator.nullsFirst(Comparator.naturalOrder()))); + final List mergedTablets = new ArrayList<>(tablets.size()); + boolean anyMerged = false; + for (final Tablet tablet : sortedTablets) { + boolean merged = false; + for (int i = 0; i < mergedTablets.size(); i++) { + final RelationalTabletWithColumnIndexMap candidate = mergedTablets.get(i); + if (canMergeRelationalTablets(candidate.tablet, candidate.columnIndexMap, tablet)) { + mergedTablets.set(i, mergeRelationalTablets(candidate, tablet)); + merged = true; + anyMerged = true; + break; + } + } + if (!merged) { + mergedTablets.add( + new RelationalTabletWithColumnIndexMap(tablet, getColumnIndexMap(tablet))); + } + } + if (anyMerged) { + consecutiveRelationalTabletMergeMissCount = 0; + } else { + consecutiveRelationalTabletMergeMissCount++; + } + return mergedTablets.stream() + .map(relationalTablet -> relationalTablet.tablet) + .collect(Collectors.toList()); + } + + private boolean canMergeRelationalTablets( + final Tablet left, final Map leftColumnIndexMap, final Tablet right) { + if (!Objects.equals(left.getTableName(), right.getTableName())) { + return false; + } + final List rightSchemas = right.getSchemas(); + int duplicatedColumnCount = 0; + for (int rightIndex = 0; rightIndex < rightSchemas.size(); rightIndex++) { + final IMeasurementSchema rightSchema = rightSchemas.get(rightIndex); + final Integer leftIndex = leftColumnIndexMap.get(rightSchema.getMeasurementName()); + if (leftIndex == null) { + continue; + } + if (left.getSchemas().get(leftIndex).getType() != rightSchema.getType() + || left.getColumnTypes().get(leftIndex) != right.getColumnTypes().get(rightIndex)) { + return false; + } + duplicatedColumnCount++; + } + return (double) duplicatedColumnCount / Math.max(left.getSchemas().size(), rightSchemas.size()) + > 0.5; + } + + private Map getColumnIndexMap(final Tablet tablet) { + final Map columnIndexMap = new HashMap<>(tablet.getSchemas().size()); + for (int i = 0; i < tablet.getSchemas().size(); i++) { + columnIndexMap.put(tablet.getSchemas().get(i).getMeasurementName(), i); + } + return columnIndexMap; + } + + private Map getColumnIndexMap(final List measurements) { + final Map columnIndexMap = new HashMap<>(measurements.size()); + for (int i = 0; i < measurements.size(); i++) { + columnIndexMap.put(measurements.get(i), i); + } + return columnIndexMap; + } + + private RelationalTabletWithColumnIndexMap mergeRelationalTablets( + final RelationalTabletWithColumnIndexMap left, final Tablet right) { + final List measurements = + new ArrayList<>(left.tablet.getSchemas().size() + right.getSchemas().size()); + final List dataTypes = + new ArrayList<>(left.tablet.getSchemas().size() + right.getSchemas().size()); + final List columnTypes = + new ArrayList<>(left.tablet.getSchemas().size() + right.getSchemas().size()); + for (int i = 0; i < left.tablet.getSchemas().size(); i++) { + measurements.add(left.tablet.getSchemas().get(i).getMeasurementName()); + dataTypes.add(left.tablet.getSchemas().get(i).getType()); + columnTypes.add(left.tablet.getColumnTypes().get(i)); + } + for (int i = 0; i < right.getSchemas().size(); i++) { + final IMeasurementSchema rightSchema = right.getSchemas().get(i); + if (!left.columnIndexMap.containsKey(rightSchema.getMeasurementName())) { + measurements.add(rightSchema.getMeasurementName()); + dataTypes.add(rightSchema.getType()); + columnTypes.add(right.getColumnTypes().get(i)); + } + } + + final int rowCount = left.tablet.getRowSize() + right.getRowSize(); + final Tablet merged = + new Tablet(left.tablet.getTableName(), measurements, dataTypes, columnTypes, rowCount); + final long[] timestamps = new long[rowCount]; + final Object[] values = new Object[measurements.size()]; + final BitMap[] bitMaps = new BitMap[measurements.size()]; + for (int i = 0; i < measurements.size(); i++) { + values[i] = createValueColumn(dataTypes.get(i), rowCount); + } + + copyTabletColumns(left.tablet, 0, measurements, values, bitMaps, timestamps); + copyTabletColumns(right, left.tablet.getRowSize(), measurements, values, bitMaps, timestamps); + merged.setTimestamps(timestamps); + merged.setValues(values); + merged.setBitMaps(bitMaps); + merged.setRowSize(rowCount); + return new RelationalTabletWithColumnIndexMap(merged, getColumnIndexMap(measurements)); + } + + private void copyTabletColumns( + final Tablet source, + final int targetRowOffset, + final List targetMeasurements, + final Object[] targetValues, + final BitMap[] targetBitMaps, + final long[] targetTimestamps) { + final int rowSize = source.getRowSize(); + System.arraycopy(source.getTimestamps(), 0, targetTimestamps, targetRowOffset, rowSize); + final Object[] sourceValues = source.getValues(); + final BitMap[] sourceBitMaps = source.getBitMaps(); + final Map sourceColumnIndexMap = getColumnIndexMap(source); + for (int column = 0; column < targetMeasurements.size(); column++) { + final Integer sourceColumn = sourceColumnIndexMap.get(targetMeasurements.get(column)); + if (sourceColumn == null) { + markNulls(targetBitMaps, column, targetTimestamps.length, targetRowOffset, rowSize); + continue; + } + System.arraycopy( + sourceValues[sourceColumn], 0, targetValues[column], targetRowOffset, rowSize); + if (sourceBitMaps != null && sourceBitMaps[sourceColumn] != null) { + for (int row = 0; row < rowSize; row++) { + if (sourceBitMaps[sourceColumn].isMarked(row)) { + markNull(targetBitMaps, column, targetTimestamps.length, targetRowOffset + row); + } + } + } + } + } + + private void markNulls( + final BitMap[] targetBitMaps, + final int column, + final int rowCount, + final int targetRowOffset, + final int nullCount) { + if (targetBitMaps[column] == null) { + targetBitMaps[column] = new BitMap(rowCount); + } + for (int row = 0; row < nullCount; row++) { + targetBitMaps[column].mark(targetRowOffset + row); + } + } + + private void markNull( + final BitMap[] targetBitMaps, final int column, final int rowCount, final int targetRow) { + if (targetBitMaps[column] == null) { + targetBitMaps[column] = new BitMap(rowCount); + } + targetBitMaps[column].mark(targetRow); + } + + private Object createValueColumn(final TSDataType dataType, final int rowCount) { + switch (dataType) { + case BOOLEAN: + return new boolean[rowCount]; + case INT32: + return new int[rowCount]; + case INT64: + case TIMESTAMP: + return new long[rowCount]; + case FLOAT: + return new float[rowCount]; + case DOUBLE: + return new double[rowCount]; + case DATE: + return new LocalDate[rowCount]; + case TEXT: + case STRING: + case BLOB: + case OBJECT: + return new Binary[rowCount]; + default: + throw new UnSupportedDataTypeException( + String.format( + SessionMessages.EXCEPTION_DATA_TYPE_ARG_NOT_SUPPORTED_31213160, dataType)); + } + } + + private static class RelationalTabletWithColumnIndexMap { + + private final Tablet tablet; + private final Map columnIndexMap; + + private RelationalTabletWithColumnIndexMap( + final Tablet tablet, final Map columnIndexMap) { + this.tablet = tablet; + this.columnIndexMap = columnIndexMap; + } + } + private void insertRelationalTabletWithLeaderCache(Tablet tablet) throws IoTDBConnectionException, StatementExecutionException { Map relationalTabletGroup = new HashMap<>(); @@ -4414,6 +4629,11 @@ public Builder enableRecordsAutoConvertTablet(boolean enableRecordsAutoConvertTa return this; } + public Builder enableMergeTablets(boolean enableMergeTablets) { + this.enableMergeTablets = enableMergeTablets; + return this; + } + public Builder nodeUrls(List nodeUrls) { this.nodeUrls = nodeUrls; return this; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index 2dca0d351dc80..ec64af65ff56b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -167,6 +167,18 @@ public TableSessionBuilder enableRedirection(boolean enableRedirection) { return this; } + /** + * Enables or disables client-side relational tablets merge before batch insert. + * + * @param enableMergeTablets whether to enable client-side tablets merge. + * @return the current {@link TableSessionBuilder} instance. + * @defaultValue true + */ + public TableSessionBuilder enableMergeTablets(boolean enableMergeTablets) { + this.enableMergeTablets = enableMergeTablets; + return this; + } + /** * Enables or disables automatic fetching of available DataNodes. * diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 07c0bc4d7e679..3d44d4e8ef683 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -134,6 +134,7 @@ public class SessionPool implements ISessionPool { // datanode while encountering retriable errors in current DataNode private boolean enableRedirection; private boolean enableRecordsAutoConvertTablet; + private boolean enableMergeTablets; private boolean enableQueryRedirection = false; private Map deviceIdToEndpoint; @@ -536,6 +537,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); } this.enableRecordsAutoConvertTablet = builder.enableRecordsAutoConvertTablet; + this.enableMergeTablets = builder.enableMergeTablets; this.connectionTimeoutInMs = builder.connectionTimeoutInMs; this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; @@ -600,6 +602,7 @@ private Session constructNewSession() { .thriftMaxFrameSize(thriftMaxFrameSize) .enableRedirection(enableRedirection) .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .enableMergeTablets(enableMergeTablets) .version(version) .useSSL(useSSL) .trustStore(trustStore) @@ -628,6 +631,7 @@ private Session constructNewSession() { .thriftMaxFrameSize(thriftMaxFrameSize) .enableRedirection(enableRedirection) .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .enableMergeTablets(enableMergeTablets) .version(version) .useSSL(useSSL) .trustStore(trustStore) @@ -3756,6 +3760,11 @@ public Builder enableRecordsAutoConvertTablet(boolean enableRecordsAutoConvertTa return this; } + public Builder enableMergeTablets(boolean enableMergeTablets) { + this.enableMergeTablets = enableMergeTablets; + return this; + } + public Builder connectionTimeoutInMs(int connectionTimeoutInMs) { this.connectionTimeoutInMs = connectionTimeoutInMs; return this; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index ab3e1f1c6a610..087aabce150ef 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -197,6 +197,18 @@ public TableSessionPoolBuilder enableRedirection(boolean enableRedirection) { return this; } + /** + * Enables or disables client-side relational tablets merge before batch insert. + * + * @param enableMergeTablets whether to enable client-side tablets merge. + * @return the current {@link TableSessionPoolBuilder} instance. + * @defaultValue true + */ + public TableSessionPoolBuilder enableMergeTablets(boolean enableMergeTablets) { + this.enableMergeTablets = enableMergeTablets; + return this; + } + /** * Sets the connection timeout in milliseconds. * diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java index e4d2dbbb41dc8..87bbcb7a99707 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java @@ -29,6 +29,7 @@ import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateResp; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; @@ -983,6 +984,197 @@ public void testInsertAlignedTabletsSortedEnableRedirection() Assert.assertTrue(session.isEnableQueryRedirection()); } + @Test + public void testMergeRelationalTabletsWithHighDuplicatedColumns() throws Exception { + final Tablet first = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 1, + "d1", + 11L, + 1.1); + final Tablet second = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s3"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 2, + "d2", + 22L, + true); + + final List mergedTablets = + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, second)); + + assertEquals(1, mergedTablets.size()); + final Tablet mergedTablet = mergedTablets.get(0); + assertEquals("table1", mergedTablet.getTableName()); + assertEquals(2, mergedTablet.getRowSize()); + assertEquals(Arrays.asList("tag1", "s1", "s2", "s3"), getMeasurementNames(mergedTablet)); + assertEquals(1L, mergedTablet.getTimestamp(0)); + assertEquals(2L, mergedTablet.getTimestamp(1)); + assertEquals(11L, mergedTablet.getValue(0, 1)); + assertEquals(22L, mergedTablet.getValue(1, 1)); + assertEquals(1.1, (double) mergedTablet.getValue(0, 2), 0.001); + Assert.assertNull(mergedTablet.getBitMaps()[0]); + Assert.assertNull(mergedTablet.getBitMaps()[1]); + Assert.assertTrue(mergedTablet.isNull(1, 2)); + Assert.assertTrue(mergedTablet.isNull(0, 3)); + Assert.assertTrue((boolean) mergedTablet.getValue(1, 3)); + } + + @Test + public void testMergeRelationalTabletsSkipLowDuplicatedColumnsAndDifferentTables() + throws Exception { + final Tablet first = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 1, + "d1", + 11L, + 1.1); + final Tablet second = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s3", "s4"), + Arrays.asList(TSDataType.STRING, TSDataType.INT32, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 2, + "d2", + 22, + true); + final Tablet third = + createRelationalTablet( + "table2", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 3, + "d3", + 33L, + 3.3); + final Tablet fourth = + createRelationalTablet( + "table0", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 4, + "d4", + 44L, + 4.4); + + final List mergedTablets = + Whitebox.invokeMethod( + session, "mergeRelationalTablets", Arrays.asList(first, second, third, fourth)); + + assertEquals(4, mergedTablets.size()); + Assert.assertSame(fourth, mergedTablets.get(0)); + Assert.assertSame(first, mergedTablets.get(1)); + Assert.assertSame(second, mergedTablets.get(2)); + Assert.assertSame(third, mergedTablets.get(3)); + } + + @Test + public void testMergeRelationalTabletsStopAfterConsecutiveMisses() throws Exception { + final Tablet first = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 1, + "d1", + 11L, + 1.1); + final Tablet second = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s3", "s4"), + Arrays.asList(TSDataType.STRING, TSDataType.INT32, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 2, + "d2", + 22, + true); + + for (int i = 0; i < 10; i++) { + final List mergedTablets = + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, second)); + assertEquals(2, mergedTablets.size()); + } + + final Tablet mergeable = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s3"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 3, + "d3", + 33L, + false); + final List tablets = Arrays.asList(first, mergeable); + final List mergedTablets = + Whitebox.invokeMethod(session, "mergeRelationalTablets", tablets); + Assert.assertSame(tablets, mergedTablets); + } + + @Test + public void testMergeRelationalTabletsResetConsecutiveMissesAfterMerge() throws Exception { + final Tablet first = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 1, + "d1", + 11L, + 1.1); + final Tablet unmergeable = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s3", "s4"), + Arrays.asList(TSDataType.STRING, TSDataType.INT32, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 2, + "d2", + 22, + true); + final Tablet mergeable = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s3"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 3, + "d3", + 33L, + false); + + for (int i = 0; i < 9; i++) { + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, unmergeable)); + } + final List firstMergedTablets = + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, mergeable)); + assertEquals(1, firstMergedTablets.size()); + for (int i = 0; i < 9; i++) { + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, unmergeable)); + } + + final List mergedTablets = + Whitebox.invokeMethod(session, "mergeRelationalTablets", Arrays.asList(first, mergeable)); + assertEquals(1, mergedTablets.size()); + } + @Test public void testTestInsertTablet() throws IoTDBConnectionException, StatementExecutionException { List schemas = new ArrayList<>(); @@ -1219,4 +1411,35 @@ public void testTimeoutUsingBuilder() { ISession session1 = new Session.Builder().timeOut(1).build(); assertEquals(1L, session1.getQueryTimeout()); } + + @Test + public void testEnableMergeTabletsUsingBuilder() { + Session sessionWithMergeTabletsDisabled = + new Session.Builder().enableMergeTablets(false).build(); + Assert.assertFalse( + Whitebox.getInternalState(sessionWithMergeTabletsDisabled, "enableMergeTablets")); + } + + private Tablet createRelationalTablet( + final String tableName, + final List measurements, + final List dataTypes, + final List columnTypes, + final long timestamp, + final Object... values) { + final Tablet tablet = new Tablet(tableName, measurements, dataTypes, columnTypes, 1); + tablet.addTimestamp(0, timestamp); + for (int i = 0; i < values.length; i++) { + tablet.addValue(measurements.get(i), 0, values[i]); + } + return tablet; + } + + private List getMeasurementNames(final Tablet tablet) { + final List measurementNames = new ArrayList<>(); + for (final IMeasurementSchema schema : tablet.getSchemas()) { + measurementNames.add(schema.getMeasurementName()); + } + return measurementNames; + } } From 6a8337bba153fdf997c491368e6aa1106a6c7bd9 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 10:39:02 +0800 Subject: [PATCH 5/9] disable merge when cost is too high --- .../org/apache/iotdb/session/Session.java | 34 +++++++++++++++++++ .../org/apache/iotdb/session/SessionTest.java | 24 +++++++++++++ 2 files changed, 58 insertions(+) diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 8e66fa7f8e573..2edff93c32490 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -172,7 +172,11 @@ public class Session implements ISession { private static final double SAMPLE_PROPORTION = 0.05; private static final int MIN_RECORDS_SIZE = 40; private static final int MAX_CONSECUTIVE_RELATIONAL_TABLET_MERGE_MISS_COUNT = 10; + private static final int MERGE_TABLETS_PERFORMANCE_CHECK_COUNT = 10; private int consecutiveRelationalTabletMergeMissCount = 0; + private int mergeTabletsPerformanceCheckCount = 0; + private long mergeTabletsCostInNanos = 0; + private long insertTabletsCostInNanos = 0; @SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile" protected volatile Map deviceIdToEndpoint; @@ -2869,17 +2873,26 @@ public void insertRelationalTablet(Tablet tablet) */ public void insertRelationalTablets(List tablets) throws IoTDBConnectionException, StatementExecutionException { + final boolean recordMergeTabletsCost = enableMergeTablets; + long mergeTabletsCost = 0; if (enableMergeTablets) { + final long mergeTabletsStartTime = System.nanoTime(); tablets = mergeRelationalTablets(tablets); + mergeTabletsCost = System.nanoTime() - mergeTabletsStartTime; } TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false); request.setWriteToTable(true); for (Tablet tablet : tablets) { request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes())); } + final long insertTabletsStartTime = System.nanoTime(); try { getDefaultSessionConnection().insertTablets(request); } catch (RedirectException ignored) { + } finally { + if (recordMergeTabletsCost) { + recordMergeTabletsCost(mergeTabletsCost, System.nanoTime() - insertTabletsStartTime); + } } } @@ -2914,12 +2927,33 @@ private List mergeRelationalTablets(final List tablets) { consecutiveRelationalTabletMergeMissCount = 0; } else { consecutiveRelationalTabletMergeMissCount++; + if (consecutiveRelationalTabletMergeMissCount + >= MAX_CONSECUTIVE_RELATIONAL_TABLET_MERGE_MISS_COUNT) { + enableMergeTablets = false; + } } return mergedTablets.stream() .map(relationalTablet -> relationalTablet.tablet) .collect(Collectors.toList()); } + private void recordMergeTabletsCost( + final long currentMergeTabletsCostInNanos, final long currentInsertTabletsCostInNanos) { + mergeTabletsPerformanceCheckCount++; + mergeTabletsCostInNanos += currentMergeTabletsCostInNanos; + insertTabletsCostInNanos += currentInsertTabletsCostInNanos; + if (mergeTabletsPerformanceCheckCount < MERGE_TABLETS_PERFORMANCE_CHECK_COUNT) { + return; + } + if (mergeTabletsCostInNanos > insertTabletsCostInNanos / 2) { + enableMergeTablets = false; + return; + } + mergeTabletsPerformanceCheckCount = 0; + mergeTabletsCostInNanos = 0; + insertTabletsCostInNanos = 0; + } + private boolean canMergeRelationalTablets( final Tablet left, final Map leftColumnIndexMap, final Tablet right) { if (!Objects.equals(left.getTableName(), right.getTableName())) { diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java index 87bbcb7a99707..814569362268f 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java @@ -1175,6 +1175,30 @@ public void testMergeRelationalTabletsResetConsecutiveMissesAfterMerge() throws assertEquals(1, mergedTablets.size()); } + @Test + public void testMergeTabletsDisabledWhenMergeCostExceedsHalfOfInsertCost() throws Exception { + for (int i = 0; i < 9; i++) { + Whitebox.invokeMethod(session, "recordMergeTabletsCost", 11L, 20L); + Assert.assertTrue(Whitebox.getInternalState(session, "enableMergeTablets")); + } + + Whitebox.invokeMethod(session, "recordMergeTabletsCost", 11L, 20L); + + Assert.assertFalse(Whitebox.getInternalState(session, "enableMergeTablets")); + } + + @Test + public void testMergeTabletsKeepsEnabledWhenMergeCostIsNotTooHigh() throws Exception { + for (int i = 0; i < 10; i++) { + Whitebox.invokeMethod(session, "recordMergeTabletsCost", 5L, 20L); + } + + Assert.assertTrue(Whitebox.getInternalState(session, "enableMergeTablets")); + assertEquals(0, (int) Whitebox.getInternalState(session, "mergeTabletsPerformanceCheckCount")); + assertEquals(0L, (long) Whitebox.getInternalState(session, "mergeTabletsCostInNanos")); + assertEquals(0L, (long) Whitebox.getInternalState(session, "insertTabletsCostInNanos")); + } + @Test public void testTestInsertTablet() throws IoTDBConnectionException, StatementExecutionException { List schemas = new ArrayList<>(); From fcce3aa9a78bc159575bf98676e5635205ca4d3c Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 11:13:53 +0800 Subject: [PATCH 6/9] Add ignored tablet insert performance IT --- .../it/session/IoTDBSessionRelationalIT.java | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java index fc9855a98355c..18d6f906ec5c5 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java @@ -54,9 +54,12 @@ import org.junit.Assert; import org.junit.Before; import org.junit.BeforeClass; +import org.junit.Ignore; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.File; import java.io.IOException; @@ -79,6 +82,11 @@ @Category({TableLocalStandaloneIT.class, TableClusterIT.class}) public class IoTDBSessionRelationalIT { + private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBSessionRelationalIT.class); + private static final int TABLET_PERFORMANCE_ROWS_PER_TABLET = 20; + private static final int TABLET_PERFORMANCE_REPEAT_COUNT = 3; + private static final int[] TABLET_PERFORMANCE_TABLET_COUNTS = {1, 2, 4, 8, 16}; + @BeforeClass public static void classSetUp() throws Exception { EnvFactory.getEnv().initClusterEnvironment(); @@ -674,6 +682,112 @@ public void insertRelationalTabletTest() } } + @Test + @Ignore("Performance comparison test for manual execution.") + public void compareInsertTabletAndInsertTabletsPerformanceWithIncreasingTabletCount() + throws IoTDBConnectionException, StatementExecutionException { + try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { + session.executeNonQueryStatement("USE \"db1\""); + createTabletPerformanceTable(session, "single_tablet_perf"); + createTabletPerformanceTable(session, "multi_tablets_perf"); + + long expectedRows = 0; + for (int tabletCount : TABLET_PERFORMANCE_TABLET_COUNTS) { + long singleTabletCost = 0; + long multiTabletsCost = 0; + final long stageBaseTime = expectedRows; + for (int repeatIndex = 0; repeatIndex < TABLET_PERFORMANCE_REPEAT_COUNT; repeatIndex++) { + final long baseTime = + stageBaseTime + (long) tabletCount * TABLET_PERFORMANCE_ROWS_PER_TABLET * repeatIndex; + final List singleTablets = + createTabletPerformanceTablets("single_tablet_perf", tabletCount, baseTime); + final List multiTablets = + createTabletPerformanceTablets("multi_tablets_perf", tabletCount, baseTime); + + final long singleTabletStartTime = System.nanoTime(); + for (Tablet tablet : singleTablets) { + session.insert(tablet); + } + singleTabletCost += System.nanoTime() - singleTabletStartTime; + + final long multiTabletsStartTime = System.nanoTime(); + session.insertTablets(multiTablets); + multiTabletsCost += System.nanoTime() - multiTabletsStartTime; + } + + expectedRows += + (long) tabletCount + * TABLET_PERFORMANCE_REPEAT_COUNT + * TABLET_PERFORMANCE_ROWS_PER_TABLET; + assertEquals( + expectedRows, + queryTabletPerformanceRowCount(session, "select count(s1) from single_tablet_perf")); + assertEquals( + expectedRows, + queryTabletPerformanceRowCount(session, "select count(s1) from multi_tablets_perf")); + Assert.assertTrue(singleTabletCost > 0); + Assert.assertTrue(multiTabletsCost > 0); + LOGGER.info( + "Relational tablet insert performance with tabletCount={}, rowsPerTablet={}, " + + "repeatCount={}, insertTablet={} ms, insertTablets={} ms, insertTablets/" + + "insertTablet ratio={}", + tabletCount, + TABLET_PERFORMANCE_ROWS_PER_TABLET, + TABLET_PERFORMANCE_REPEAT_COUNT, + singleTabletCost / 1_000_000, + multiTabletsCost / 1_000_000, + String.format("%.3f", (double) multiTabletsCost / singleTabletCost)); + } + } + } + + private void createTabletPerformanceTable(final ITableSession session, final String tableName) + throws IoTDBConnectionException, StatementExecutionException { + session.executeNonQueryStatement( + "CREATE TABLE " + + tableName + + " (tag1 string tag, attr1 string attribute, s1 int64 field, s2 double field)"); + } + + private List createTabletPerformanceTablets( + final String tableName, final int tabletCount, final long baseTime) { + final List tablets = new ArrayList<>(tabletCount); + for (int tabletIndex = 0; tabletIndex < tabletCount; tabletIndex++) { + final Tablet tablet = + new Tablet( + tableName, + Arrays.asList("tag1", "attr1", "s1", "s2"), + Arrays.asList( + TSDataType.STRING, TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList( + ColumnCategory.TAG, + ColumnCategory.ATTRIBUTE, + ColumnCategory.FIELD, + ColumnCategory.FIELD), + TABLET_PERFORMANCE_ROWS_PER_TABLET); + for (int row = 0; row < TABLET_PERFORMANCE_ROWS_PER_TABLET; row++) { + final int rowIndex = tablet.getRowSize(); + final long timestamp = + baseTime + (long) tabletIndex * TABLET_PERFORMANCE_ROWS_PER_TABLET + row; + tablet.addTimestamp(rowIndex, timestamp); + tablet.addValue("tag1", rowIndex, "tag:" + tabletIndex); + tablet.addValue("attr1", rowIndex, "attr:" + row); + tablet.addValue("s1", rowIndex, timestamp); + tablet.addValue("s2", rowIndex, timestamp * 1.0); + } + tablets.add(tablet); + } + return tablets; + } + + private long queryTabletPerformanceRowCount(final ITableSession session, final String query) + throws IoTDBConnectionException, StatementExecutionException { + try (SessionDataSet dataSet = session.executeQueryStatement(query)) { + assertTrue(dataSet.hasNext()); + return dataSet.next().getFields().get(0).getLongV(); + } + } + @Test public void insertNoFieldTest() throws IoTDBConnectionException, StatementExecutionException { try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { From 4073a0b423782636f842ffd38f78070cc7e2abcd Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 11:45:45 +0800 Subject: [PATCH 7/9] Warm up tablet insert performance IT --- .../it/session/IoTDBSessionRelationalIT.java | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java index 18d6f906ec5c5..e24c093de1fdf 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java @@ -85,6 +85,7 @@ public class IoTDBSessionRelationalIT { private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBSessionRelationalIT.class); private static final int TABLET_PERFORMANCE_ROWS_PER_TABLET = 20; private static final int TABLET_PERFORMANCE_REPEAT_COUNT = 3; + private static final int TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT = 16; private static final int[] TABLET_PERFORMANCE_TABLET_COUNTS = {1, 2, 4, 8, 16}; @BeforeClass @@ -691,7 +692,7 @@ public void compareInsertTabletAndInsertTabletsPerformanceWithIncreasingTabletCo createTabletPerformanceTable(session, "single_tablet_perf"); createTabletPerformanceTable(session, "multi_tablets_perf"); - long expectedRows = 0; + long expectedRows = warmUpTabletPerformanceTest(session); for (int tabletCount : TABLET_PERFORMANCE_TABLET_COUNTS) { long singleTabletCost = 0; long multiTabletsCost = 0; @@ -749,6 +750,30 @@ private void createTabletPerformanceTable(final ITableSession session, final Str + " (tag1 string tag, attr1 string attribute, s1 int64 field, s2 double field)"); } + private long warmUpTabletPerformanceTest(final ITableSession session) + throws IoTDBConnectionException, StatementExecutionException { + final List singleTablets = + createTabletPerformanceTablets( + "single_tablet_perf", TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT, 0); + final List multiTablets = + createTabletPerformanceTablets( + "multi_tablets_perf", TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT, 0); + for (Tablet tablet : singleTablets) { + session.insert(tablet); + } + session.insertTablets(multiTablets); + + final long expectedRows = + (long) TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT * TABLET_PERFORMANCE_ROWS_PER_TABLET; + assertEquals( + expectedRows, + queryTabletPerformanceRowCount(session, "select count(s1) from single_tablet_perf")); + assertEquals( + expectedRows, + queryTabletPerformanceRowCount(session, "select count(s1) from multi_tablets_perf")); + return expectedRows; + } + private List createTabletPerformanceTablets( final String tableName, final int tabletCount, final long baseTime) { final List tablets = new ArrayList<>(tabletCount); From f01ee5ea484be84a29311d48124df3d7b81bd8b8 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 12:08:28 +0800 Subject: [PATCH 8/9] Stabilize tablet insert performance IT --- .../it/session/IoTDBSessionRelationalIT.java | 32 +++++++++++++------ 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java index e24c093de1fdf..6975dd3873f55 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java @@ -84,7 +84,7 @@ public class IoTDBSessionRelationalIT { private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBSessionRelationalIT.class); private static final int TABLET_PERFORMANCE_ROWS_PER_TABLET = 20; - private static final int TABLET_PERFORMANCE_REPEAT_COUNT = 3; + private static final int TABLET_PERFORMANCE_REPEAT_COUNT = 20; private static final int TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT = 16; private static final int[] TABLET_PERFORMANCE_TABLET_COUNTS = {1, 2, 4, 8, 16}; @@ -705,15 +705,13 @@ public void compareInsertTabletAndInsertTabletsPerformanceWithIncreasingTabletCo final List multiTablets = createTabletPerformanceTablets("multi_tablets_perf", tabletCount, baseTime); - final long singleTabletStartTime = System.nanoTime(); - for (Tablet tablet : singleTablets) { - session.insert(tablet); + if (repeatIndex % 2 == 0) { + singleTabletCost += insertTabletPerformanceTest(session, singleTablets); + multiTabletsCost += insertTabletsPerformanceTest(session, multiTablets); + } else { + multiTabletsCost += insertTabletsPerformanceTest(session, multiTablets); + singleTabletCost += insertTabletPerformanceTest(session, singleTablets); } - singleTabletCost += System.nanoTime() - singleTabletStartTime; - - final long multiTabletsStartTime = System.nanoTime(); - session.insertTablets(multiTablets); - multiTabletsCost += System.nanoTime() - multiTabletsStartTime; } expectedRows += @@ -774,6 +772,22 @@ private long warmUpTabletPerformanceTest(final ITableSession session) return expectedRows; } + private long insertTabletPerformanceTest(final ITableSession session, final List tablets) + throws IoTDBConnectionException, StatementExecutionException { + final long startTime = System.nanoTime(); + for (Tablet tablet : tablets) { + session.insert(tablet); + } + return System.nanoTime() - startTime; + } + + private long insertTabletsPerformanceTest(final ITableSession session, final List tablets) + throws IoTDBConnectionException, StatementExecutionException { + final long startTime = System.nanoTime(); + session.insertTablets(tablets); + return System.nanoTime() - startTime; + } + private List createTabletPerformanceTablets( final String tableName, final int tabletCount, final long baseTime) { final List tablets = new ArrayList<>(tabletCount); From e7eefcdc1c5a3b94a3ce40dd24c19dd35679c8d1 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Tue, 7 Jul 2026 14:45:59 +0800 Subject: [PATCH 9/9] Address relational insertTablets review comments --- .../it/session/IoTDBSessionRelationalIT.java | 15 ++ .../org/apache/iotdb/session/Session.java | 224 +++++++++++++++++- .../org/apache/iotdb/session/SessionTest.java | 68 ++++++ .../relational/sql/ast/InsertTablets.java | 2 +- 4 files changed, 299 insertions(+), 10 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java index 6975dd3873f55..6771464be840c 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java @@ -683,6 +683,21 @@ public void insertRelationalTabletTest() } } + @Test + public void insertTabletsAutoCreateTableTest() + throws IoTDBConnectionException, StatementExecutionException { + try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { + session.executeNonQueryStatement("USE \"db1\""); + final List tablets = createTabletPerformanceTablets("auto_create_tablets", 2, 0); + + session.insertTablets(tablets); + + assertEquals( + (long) TABLET_PERFORMANCE_ROWS_PER_TABLET * 2, + queryTabletPerformanceRowCount(session, "select count(s1) from auto_create_tablets")); + } + } + @Test @Ignore("Performance comparison test for manual execution.") public void compareInsertTabletAndInsertTabletsPerformanceWithIncreasingTabletCount() diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 2edff93c32490..e39b5e8d37732 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -92,6 +92,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; @@ -2880,18 +2881,199 @@ public void insertRelationalTablets(List tablets) tablets = mergeRelationalTablets(tablets); mergeTabletsCost = System.nanoTime() - mergeTabletsStartTime; } - TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false); - request.setWriteToTable(true); - for (Tablet tablet : tablets) { - request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes())); + if (tablets.isEmpty()) { + throw new BatchExecutionException(SessionMessages.NO_TABLET_INSERTING); } final long insertTabletsStartTime = System.nanoTime(); + if (enableRedirection) { + insertRelationalTabletsWithLeaderCache(tablets); + } else { + TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false); + request.setWriteToTable(true); + for (Tablet tablet : tablets) { + request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes())); + } + try { + getDefaultSessionConnection().insertTablets(request); + } catch (RedirectException ignored) { + } + } + if (recordMergeTabletsCost) { + recordMergeTabletsCost(mergeTabletsCost, System.nanoTime() - insertTabletsStartTime); + } + } + + private void insertRelationalTabletsWithLeaderCache(final List tablets) + throws IoTDBConnectionException, StatementExecutionException { + final Map tabletGroup = new HashMap<>(); + for (final Tablet tablet : tablets) { + addRelationalTabletToGroup(tabletGroup, tablet); + } + if (tabletGroup.size() == 1) { + insertRelationalTabletsOnce(tabletGroup); + } else { + insertRelationalTabletsByGroup(tabletGroup); + } + } + + private void addRelationalTabletToGroup( + final Map tabletGroup, final Tablet tablet) + throws IoTDBConnectionException { + if (SessionUtils.isTabletContainsSingleDevice(tablet)) { + final SessionConnection connection = getTableModelSessionConnection(tablet.getDeviceID(0)); + updateRelationalTabletsReq(tabletGroup, connection, tablet); + return; + } + for (int row = 0; row < tablet.getRowSize(); row++) { + final IDeviceID deviceID = tablet.getDeviceID(row); + final SessionConnection connection = getTableModelSessionConnection(deviceID); + final Tablet subTablet = + tabletGroup + .computeIfAbsent(connection, ignored -> new RelationalTabletsReq()) + .tabletByDevice + .computeIfAbsent(deviceID, ignored -> createEmptyRelationalTabletLike(tablet)); + for (int column = 0; column < subTablet.getSchemas().size(); column++) { + subTablet.addValue( + subTablet.getSchemas().get(column).getMeasurementName(), + subTablet.getRowSize(), + tablet.getValue(row, column)); + } + subTablet.addTimestamp(subTablet.getRowSize(), tablet.getTimestamp(row)); + } + } + + private SessionConnection getTableModelSessionConnection(final IDeviceID deviceID) + throws IoTDBConnectionException { + return tableModelDeviceIdToEndpoint == null || tableModelDeviceIdToEndpoint.isEmpty() + ? getDefaultSessionConnection() + : getSessionConnection(deviceID); + } + + private Tablet createEmptyRelationalTabletLike(final Tablet tablet) { + final List measurements = new ArrayList<>(tablet.getSchemas().size()); + final List dataTypes = new ArrayList<>(tablet.getSchemas().size()); + tablet + .getSchemas() + .forEach( + schema -> { + measurements.add(schema.getMeasurementName()); + dataTypes.add(schema.getType()); + }); + return new Tablet( + tablet.getTableName(), + measurements, + dataTypes, + tablet.getColumnTypes(), + tablet.getRowSize()); + } + + private void updateRelationalTabletsReq( + final Map tabletGroup, + final SessionConnection connection, + final Tablet tablet) { + tabletGroup + .computeIfAbsent(connection, ignored -> new RelationalTabletsReq()) + .addTablet(tablet); + } + + private void buildRelationalTabletsRequests( + final Map tabletGroup) { + for (final RelationalTabletsReq relationalTabletsReq : tabletGroup.values()) { + relationalTabletsReq.buildRequest(); + } + } + + private void insertRelationalTabletsOnce( + final Map tabletGroup) + throws IoTDBConnectionException, StatementExecutionException { + buildRelationalTabletsRequests(tabletGroup); + final Map.Entry entry = + tabletGroup.entrySet().iterator().next(); + final SessionConnection connection = entry.getKey(); + final RelationalTabletsReq relationalTabletsReq = entry.getValue(); try { - getDefaultSessionConnection().insertTablets(request); - } catch (RedirectException ignored) { - } finally { - if (recordMergeTabletsCost) { - recordMergeTabletsCost(mergeTabletsCost, System.nanoTime() - insertTabletsStartTime); + connection.insertTablets(relationalTabletsReq.request); + } catch (RedirectException e) { + handleRelationalTabletsRedirection(e, relationalTabletsReq.deviceIDs); + } catch (IoTDBConnectionException e) { + if (endPointToSessionConnection != null && endPointToSessionConnection.size() > 1) { + removeBrokenSessionConnection(connection); + try { + getDefaultSessionConnection().insertTablets(relationalTabletsReq.request); + } catch (RedirectException ignored) { + } + } else { + throw e; + } + } + } + + @SuppressWarnings({ + "squid:S3776" + }) // ignore Cognitive Complexity of methods should not be too high + private void insertRelationalTabletsByGroup( + final Map tabletGroup) + throws IoTDBConnectionException, StatementExecutionException { + buildRelationalTabletsRequests(tabletGroup); + final List> completableFutures = + tabletGroup.entrySet().stream() + .map( + entry -> { + final SessionConnection connection = entry.getKey(); + final RelationalTabletsReq relationalTabletsReq = entry.getValue(); + return CompletableFuture.runAsync( + () -> { + try { + connection.insertTablets(relationalTabletsReq.request); + } catch (RedirectException e) { + handleRelationalTabletsRedirection(e, relationalTabletsReq.deviceIDs); + } catch (StatementExecutionException e) { + throw new CompletionException(e); + } catch (IoTDBConnectionException e) { + removeBrokenSessionConnection(connection); + try { + getDefaultSessionConnection() + .insertTablets(relationalTabletsReq.request); + } catch (IoTDBConnectionException | StatementExecutionException ex) { + throw new CompletionException(ex); + } catch (RedirectException ignored) { + } + } + }, + OPERATION_EXECUTOR); + }) + .collect(Collectors.toList()); + + final StringBuilder errMsgBuilder = new StringBuilder(); + for (final CompletableFuture completableFuture : completableFutures) { + try { + completableFuture.join(); + } catch (CompletionException completionException) { + final Throwable cause = completionException.getCause(); + logger.error(SessionMessages.MEET_ERROR_WHEN_ASYNC_INSERT, cause); + if (cause instanceof IoTDBConnectionException) { + throw (IoTDBConnectionException) cause; + } + if (errMsgBuilder.length() > 0) { + errMsgBuilder.append(";"); + } + errMsgBuilder.append(cause.getMessage()); + } + } + if (errMsgBuilder.length() > 0) { + throw new StatementExecutionException(errMsgBuilder.toString()); + } + } + + private void handleRelationalTabletsRedirection( + final RedirectException redirectException, final List deviceIDs) { + final List endPointList = redirectException.getEndPointList(); + if (endPointList == null) { + return; + } + for (int i = 0; i < endPointList.size() && i < deviceIDs.size(); i++) { + if (endPointList.get(i) != null) { + handleRedirection(deviceIDs.get(i), endPointList.get(i)); } } } @@ -3125,6 +3307,30 @@ private RelationalTabletWithColumnIndexMap( } } + private class RelationalTabletsReq { + + private final TSInsertTabletsReq request = new TSInsertTabletsReq(); + private final List deviceIDs = new ArrayList<>(); + private final List tablets = new ArrayList<>(); + private final Map tabletByDevice = new LinkedHashMap<>(); + + private void addTablet(final Tablet tablet) { + tablets.add(tablet); + } + + private void buildRequest() { + request.setWriteToTable(true); + tablets.forEach(this::addTabletToRequest); + tabletByDevice.values().forEach(this::addTabletToRequest); + } + + private void addTabletToRequest(final Tablet tablet) { + updateTSInsertTabletsReq(request, tablet, false, false); + request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes())); + deviceIDs.add(tablet.getDeviceID(0)); + } + } + private void insertRelationalTabletWithLeaderCache(Tablet tablet) throws IoTDBConnectionException, StatementExecutionException { Map relationalTabletGroup = new HashMap<>(); diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java index 814569362268f..0b317311ce7f5 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java @@ -26,11 +26,13 @@ import org.apache.iotdb.isession.util.Version; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq; import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateResp; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.utils.Binary; @@ -1199,6 +1201,72 @@ public void testMergeTabletsKeepsEnabledWhenMergeCostIsNotTooHigh() throws Excep assertEquals(0L, (long) Whitebox.getInternalState(session, "insertTabletsCostInNanos")); } + @Test + public void testMergeTabletsCostNotRecordedWhenInsertFails() throws Exception { + Mockito.doThrow(new StatementExecutionException("expected")) + .when(sessionConnection) + .insertTablets(any(TSInsertTabletsReq.class)); + final Tablet first = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s2"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 1L, + "tag1", + 11L, + 1.1); + final Tablet second = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1", "s3"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.BOOLEAN), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD), + 2L, + "tag2", + 22L, + true); + + try { + ((Session) session).insertRelationalTablets(Arrays.asList(first, second)); + fail("Exception expected"); + } catch (StatementExecutionException e) { + assertEquals("expected", e.getMessage()); + } + + assertEquals(0, (int) Whitebox.getInternalState(session, "mergeTabletsPerformanceCheckCount")); + assertEquals(0L, (long) Whitebox.getInternalState(session, "mergeTabletsCostInNanos")); + assertEquals(0L, (long) Whitebox.getInternalState(session, "insertTabletsCostInNanos")); + } + + @Test + public void testInsertRelationalTabletsUseTableModelLeaderCache() throws Exception { + final Tablet tablet = + createRelationalTablet( + "table1", + Arrays.asList("tag1", "s1"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1L, + "tag1", + 11L); + final SessionConnection redirectedSessionConnection = Mockito.mock(SessionConnection.class); + final TEndPoint endPoint = new TEndPoint("127.0.0.2", 6667); + final Map tableModelDeviceIdToEndpoint = new HashMap<>(); + tableModelDeviceIdToEndpoint.put(tablet.getDeviceID(0), endPoint); + final Map endPointToSessionConnection = new HashMap<>(); + endPointToSessionConnection.put(endPoint, redirectedSessionConnection); + Whitebox.setInternalState(session, "enableRedirection", true); + Whitebox.setInternalState( + session, "tableModelDeviceIdToEndpoint", tableModelDeviceIdToEndpoint); + Whitebox.setInternalState(session, "endPointToSessionConnection", endPointToSessionConnection); + + ((Session) session).insertRelationalTablets(Collections.singletonList(tablet)); + + Mockito.verify(redirectedSessionConnection).insertTablets(any(TSInsertTabletsReq.class)); + Mockito.verify(sessionConnection, Mockito.never()).insertTablets(any(TSInsertTabletsReq.class)); + } + @Test public void testTestInsertTablet() throws IoTDBConnectionException, StatementExecutionException { List schemas = new ArrayList<>(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java index fbccd7800d028..d3f64db6aba6d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java @@ -76,7 +76,7 @@ public void validateTableSchema(Metadata metadata, MPPQueryContext context) { for (InsertTabletStatement insertTabletStatement : getInnerTreeStatement().getInsertTabletStatementList()) { final String database = AnalyzeUtils.getDatabaseName(insertTabletStatement, context); - super.validateTableSchema(metadata, context, insertTabletStatement, database, false); + super.validateTableSchema(metadata, context, insertTabletStatement, database, true); } }