Optimize relational tablet batch inserts

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,12 @@
 @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 = 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};
+
   @BeforeClass
   public static void classSetUp() throws Exception {
     EnvFactory.getEnv().initClusterEnvironment();
@@ -674,6 +683,165 @@ public void insertRelationalTabletTest()
     }
   }
 
+  @Test
+  public void insertTabletsAutoCreateTableTest()
+      throws IoTDBConnectionException, StatementExecutionException {
+    try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) {
+      session.executeNonQueryStatement("USE \"db1\"");
+      final List<Tablet> 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()
+      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 = warmUpTabletPerformanceTest(session);
+      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<Tablet> singleTablets =
+              createTabletPerformanceTablets("single_tablet_perf", tabletCount, baseTime);
+          final List<Tablet> multiTablets =
+              createTabletPerformanceTablets("multi_tablets_perf", tabletCount, baseTime);
+
+          if (repeatIndex % 2 == 0) {
+            singleTabletCost += insertTabletPerformanceTest(session, singleTablets);
+            multiTabletsCost += insertTabletsPerformanceTest(session, multiTablets);
+          } else {
+            multiTabletsCost += insertTabletsPerformanceTest(session, multiTablets);
+            singleTabletCost += insertTabletPerformanceTest(session, singleTablets);
+          }
+        }
+
+        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 long warmUpTabletPerformanceTest(final ITableSession session)
+      throws IoTDBConnectionException, StatementExecutionException {
+    final List<Tablet> singleTablets =
+        createTabletPerformanceTablets(
+            "single_tablet_perf", TABLET_PERFORMANCE_WARM_UP_TABLET_COUNT, 0);
+    final List<Tablet> 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 long insertTabletPerformanceTest(final ITableSession session, final List<Tablet> 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<Tablet> tablets)
+      throws IoTDBConnectionException, StatementExecutionException {
+    final long startTime = System.nanoTime();
+    session.insertTablets(tablets);
+    return System.nanoTime() - startTime;
+  }
+
+  private List<Tablet> createTabletPerformanceTablets(
+      final String tableName, final int tabletCount, final long baseTime) {
+    final List<Tablet> 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()) {

integration-test/src/test/java/org/apache/iotdb/relational/it/session/pool/IoTDBInsertTableSessionPoolIT.java

@@ -250,6 +250,187 @@ 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());
+      }
+    }
+  }
+
+  @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<IMeasurementSchema> 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<ColumnCategory> 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<IMeasurementSchema> schemas = new ArrayList<>();
+    schemas.add(new MeasurementSchema("tag2", TSDataType.STRING));
+    schemas.add(new MeasurementSchema("s2", TSDataType.DOUBLE));
+    final List<ColumnCategory> 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<IMeasurementSchema> 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<ColumnCategory> 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());
     }
   }
 }