Skip to content

Commit e7eefcd

Browse files
jt2594838jt2594838
committed
Address relational insertTablets review comments
1 parent f01ee5e commit e7eefcd

4 files changed

Lines changed: 299 additions & 10 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,21 @@ public void insertRelationalTabletTest()
683683
}
684684
}
685685

686+
@Test
687+
public void insertTabletsAutoCreateTableTest()
688+
throws IoTDBConnectionException, StatementExecutionException {
689+
try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) {
690+
session.executeNonQueryStatement("USE \"db1\"");
691+
final List<Tablet> tablets = createTabletPerformanceTablets("auto_create_tablets", 2, 0);
692+
693+
session.insertTablets(tablets);
694+
695+
assertEquals(
696+
(long) TABLET_PERFORMANCE_ROWS_PER_TABLET * 2,
697+
queryTabletPerformanceRowCount(session, "select count(s1) from auto_create_tablets"));
698+
}
699+
}
700+
686701
@Test
687702
@Ignore("Performance comparison test for manual execution.")
688703
public void compareInsertTabletAndInsertTabletsPerformanceWithIncreasingTabletCount()

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java‎

Lines changed: 215 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@
9292
import java.util.HashMap;
9393
import java.util.HashSet;
9494
import java.util.Iterator;
95+
import java.util.LinkedHashMap;
9596
import java.util.List;
9697
import java.util.Map;
9798
import java.util.Map.Entry;
@@ -2880,18 +2881,199 @@ public void insertRelationalTablets(List<Tablet> tablets)
28802881
tablets = mergeRelationalTablets(tablets);
28812882
mergeTabletsCost = System.nanoTime() - mergeTabletsStartTime;
28822883
}
2883-
TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false);
2884-
request.setWriteToTable(true);
2885-
for (Tablet tablet : tablets) {
2886-
request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes()));
2884+
if (tablets.isEmpty()) {
2885+
throw new BatchExecutionException(SessionMessages.NO_TABLET_INSERTING);
28872886
}
28882887
final long insertTabletsStartTime = System.nanoTime();
2888+
if (enableRedirection) {
2889+
insertRelationalTabletsWithLeaderCache(tablets);
2890+
} else {
2891+
TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, false, false);
2892+
request.setWriteToTable(true);
2893+
for (Tablet tablet : tablets) {
2894+
request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes()));
2895+
}
2896+
try {
2897+
getDefaultSessionConnection().insertTablets(request);
2898+
} catch (RedirectException ignored) {
2899+
}
2900+
}
2901+
if (recordMergeTabletsCost) {
2902+
recordMergeTabletsCost(mergeTabletsCost, System.nanoTime() - insertTabletsStartTime);
2903+
}
2904+
}
2905+
2906+
private void insertRelationalTabletsWithLeaderCache(final List<Tablet> tablets)
2907+
throws IoTDBConnectionException, StatementExecutionException {
2908+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup = new HashMap<>();
2909+
for (final Tablet tablet : tablets) {
2910+
addRelationalTabletToGroup(tabletGroup, tablet);
2911+
}
2912+
if (tabletGroup.size() == 1) {
2913+
insertRelationalTabletsOnce(tabletGroup);
2914+
} else {
2915+
insertRelationalTabletsByGroup(tabletGroup);
2916+
}
2917+
}
2918+
2919+
private void addRelationalTabletToGroup(
2920+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup, final Tablet tablet)
2921+
throws IoTDBConnectionException {
2922+
if (SessionUtils.isTabletContainsSingleDevice(tablet)) {
2923+
final SessionConnection connection = getTableModelSessionConnection(tablet.getDeviceID(0));
2924+
updateRelationalTabletsReq(tabletGroup, connection, tablet);
2925+
return;
2926+
}
2927+
for (int row = 0; row < tablet.getRowSize(); row++) {
2928+
final IDeviceID deviceID = tablet.getDeviceID(row);
2929+
final SessionConnection connection = getTableModelSessionConnection(deviceID);
2930+
final Tablet subTablet =
2931+
tabletGroup
2932+
.computeIfAbsent(connection, ignored -> new RelationalTabletsReq())
2933+
.tabletByDevice
2934+
.computeIfAbsent(deviceID, ignored -> createEmptyRelationalTabletLike(tablet));
2935+
for (int column = 0; column < subTablet.getSchemas().size(); column++) {
2936+
subTablet.addValue(
2937+
subTablet.getSchemas().get(column).getMeasurementName(),
2938+
subTablet.getRowSize(),
2939+
tablet.getValue(row, column));
2940+
}
2941+
subTablet.addTimestamp(subTablet.getRowSize(), tablet.getTimestamp(row));
2942+
}
2943+
}
2944+
2945+
private SessionConnection getTableModelSessionConnection(final IDeviceID deviceID)
2946+
throws IoTDBConnectionException {
2947+
return tableModelDeviceIdToEndpoint == null || tableModelDeviceIdToEndpoint.isEmpty()
2948+
? getDefaultSessionConnection()
2949+
: getSessionConnection(deviceID);
2950+
}
2951+
2952+
private Tablet createEmptyRelationalTabletLike(final Tablet tablet) {
2953+
final List<String> measurements = new ArrayList<>(tablet.getSchemas().size());
2954+
final List<TSDataType> dataTypes = new ArrayList<>(tablet.getSchemas().size());
2955+
tablet
2956+
.getSchemas()
2957+
.forEach(
2958+
schema -> {
2959+
measurements.add(schema.getMeasurementName());
2960+
dataTypes.add(schema.getType());
2961+
});
2962+
return new Tablet(
2963+
tablet.getTableName(),
2964+
measurements,
2965+
dataTypes,
2966+
tablet.getColumnTypes(),
2967+
tablet.getRowSize());
2968+
}
2969+
2970+
private void updateRelationalTabletsReq(
2971+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup,
2972+
final SessionConnection connection,
2973+
final Tablet tablet) {
2974+
tabletGroup
2975+
.computeIfAbsent(connection, ignored -> new RelationalTabletsReq())
2976+
.addTablet(tablet);
2977+
}
2978+
2979+
private void buildRelationalTabletsRequests(
2980+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup) {
2981+
for (final RelationalTabletsReq relationalTabletsReq : tabletGroup.values()) {
2982+
relationalTabletsReq.buildRequest();
2983+
}
2984+
}
2985+
2986+
private void insertRelationalTabletsOnce(
2987+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup)
2988+
throws IoTDBConnectionException, StatementExecutionException {
2989+
buildRelationalTabletsRequests(tabletGroup);
2990+
final Map.Entry<SessionConnection, RelationalTabletsReq> entry =
2991+
tabletGroup.entrySet().iterator().next();
2992+
final SessionConnection connection = entry.getKey();
2993+
final RelationalTabletsReq relationalTabletsReq = entry.getValue();
28892994
try {
2890-
getDefaultSessionConnection().insertTablets(request);
2891-
} catch (RedirectException ignored) {
2892-
} finally {
2893-
if (recordMergeTabletsCost) {
2894-
recordMergeTabletsCost(mergeTabletsCost, System.nanoTime() - insertTabletsStartTime);
2995+
connection.insertTablets(relationalTabletsReq.request);
2996+
} catch (RedirectException e) {
2997+
handleRelationalTabletsRedirection(e, relationalTabletsReq.deviceIDs);
2998+
} catch (IoTDBConnectionException e) {
2999+
if (endPointToSessionConnection != null && endPointToSessionConnection.size() > 1) {
3000+
removeBrokenSessionConnection(connection);
3001+
try {
3002+
getDefaultSessionConnection().insertTablets(relationalTabletsReq.request);
3003+
} catch (RedirectException ignored) {
3004+
}
3005+
} else {
3006+
throw e;
3007+
}
3008+
}
3009+
}
3010+
3011+
@SuppressWarnings({
3012+
"squid:S3776"
3013+
}) // ignore Cognitive Complexity of methods should not be too high
3014+
private void insertRelationalTabletsByGroup(
3015+
final Map<SessionConnection, RelationalTabletsReq> tabletGroup)
3016+
throws IoTDBConnectionException, StatementExecutionException {
3017+
buildRelationalTabletsRequests(tabletGroup);
3018+
final List<CompletableFuture<Void>> completableFutures =
3019+
tabletGroup.entrySet().stream()
3020+
.map(
3021+
entry -> {
3022+
final SessionConnection connection = entry.getKey();
3023+
final RelationalTabletsReq relationalTabletsReq = entry.getValue();
3024+
return CompletableFuture.runAsync(
3025+
() -> {
3026+
try {
3027+
connection.insertTablets(relationalTabletsReq.request);
3028+
} catch (RedirectException e) {
3029+
handleRelationalTabletsRedirection(e, relationalTabletsReq.deviceIDs);
3030+
} catch (StatementExecutionException e) {
3031+
throw new CompletionException(e);
3032+
} catch (IoTDBConnectionException e) {
3033+
removeBrokenSessionConnection(connection);
3034+
try {
3035+
getDefaultSessionConnection()
3036+
.insertTablets(relationalTabletsReq.request);
3037+
} catch (IoTDBConnectionException | StatementExecutionException ex) {
3038+
throw new CompletionException(ex);
3039+
} catch (RedirectException ignored) {
3040+
}
3041+
}
3042+
},
3043+
OPERATION_EXECUTOR);
3044+
})
3045+
.collect(Collectors.toList());
3046+
3047+
final StringBuilder errMsgBuilder = new StringBuilder();
3048+
for (final CompletableFuture<Void> completableFuture : completableFutures) {
3049+
try {
3050+
completableFuture.join();
3051+
} catch (CompletionException completionException) {
3052+
final Throwable cause = completionException.getCause();
3053+
logger.error(SessionMessages.MEET_ERROR_WHEN_ASYNC_INSERT, cause);
3054+
if (cause instanceof IoTDBConnectionException) {
3055+
throw (IoTDBConnectionException) cause;
3056+
}
3057+
if (errMsgBuilder.length() > 0) {
3058+
errMsgBuilder.append(";");
3059+
}
3060+
errMsgBuilder.append(cause.getMessage());
3061+
}
3062+
}
3063+
if (errMsgBuilder.length() > 0) {
3064+
throw new StatementExecutionException(errMsgBuilder.toString());
3065+
}
3066+
}
3067+
3068+
private void handleRelationalTabletsRedirection(
3069+
final RedirectException redirectException, final List<IDeviceID> deviceIDs) {
3070+
final List<TEndPoint> endPointList = redirectException.getEndPointList();
3071+
if (endPointList == null) {
3072+
return;
3073+
}
3074+
for (int i = 0; i < endPointList.size() && i < deviceIDs.size(); i++) {
3075+
if (endPointList.get(i) != null) {
3076+
handleRedirection(deviceIDs.get(i), endPointList.get(i));
28953077
}
28963078
}
28973079
}
@@ -3125,6 +3307,30 @@ private RelationalTabletWithColumnIndexMap(
31253307
}
31263308
}
31273309

3310+
private class RelationalTabletsReq {
3311+
3312+
private final TSInsertTabletsReq request = new TSInsertTabletsReq();
3313+
private final List<IDeviceID> deviceIDs = new ArrayList<>();
3314+
private final List<Tablet> tablets = new ArrayList<>();
3315+
private final Map<IDeviceID, Tablet> tabletByDevice = new LinkedHashMap<>();
3316+
3317+
private void addTablet(final Tablet tablet) {
3318+
tablets.add(tablet);
3319+
}
3320+
3321+
private void buildRequest() {
3322+
request.setWriteToTable(true);
3323+
tablets.forEach(this::addTabletToRequest);
3324+
tabletByDevice.values().forEach(this::addTabletToRequest);
3325+
}
3326+
3327+
private void addTabletToRequest(final Tablet tablet) {
3328+
updateTSInsertTabletsReq(request, tablet, false, false);
3329+
request.addToColumnCategoriesList(toEnumOrdinalsAsBytes(tablet.getColumnTypes()));
3330+
deviceIDs.add(tablet.getDeviceID(0));
3331+
}
3332+
}
3333+
31283334
private void insertRelationalTabletWithLeaderCache(Tablet tablet)
31293335
throws IoTDBConnectionException, StatementExecutionException {
31303336
Map<SessionConnection, Tablet> relationalTabletGroup = new HashMap<>();

‎iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,13 @@
2626
import org.apache.iotdb.isession.util.Version;
2727
import org.apache.iotdb.rpc.IoTDBConnectionException;
2828
import org.apache.iotdb.rpc.StatementExecutionException;
29+
import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq;
2930
import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateResp;
3031

3132
import org.apache.tsfile.common.conf.TSFileConfig;
3233
import org.apache.tsfile.enums.ColumnCategory;
3334
import org.apache.tsfile.enums.TSDataType;
35+
import org.apache.tsfile.file.metadata.IDeviceID;
3436
import org.apache.tsfile.file.metadata.enums.CompressionType;
3537
import org.apache.tsfile.file.metadata.enums.TSEncoding;
3638
import org.apache.tsfile.utils.Binary;
@@ -1199,6 +1201,72 @@ public void testMergeTabletsKeepsEnabledWhenMergeCostIsNotTooHigh() throws Excep
11991201
assertEquals(0L, (long) Whitebox.getInternalState(session, "insertTabletsCostInNanos"));
12001202
}
12011203

1204+
@Test
1205+
public void testMergeTabletsCostNotRecordedWhenInsertFails() throws Exception {
1206+
Mockito.doThrow(new StatementExecutionException("expected"))
1207+
.when(sessionConnection)
1208+
.insertTablets(any(TSInsertTabletsReq.class));
1209+
final Tablet first =
1210+
createRelationalTablet(
1211+
"table1",
1212+
Arrays.asList("tag1", "s1", "s2"),
1213+
Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.DOUBLE),
1214+
Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD),
1215+
1L,
1216+
"tag1",
1217+
11L,
1218+
1.1);
1219+
final Tablet second =
1220+
createRelationalTablet(
1221+
"table1",
1222+
Arrays.asList("tag1", "s1", "s3"),
1223+
Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.BOOLEAN),
1224+
Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD),
1225+
2L,
1226+
"tag2",
1227+
22L,
1228+
true);
1229+
1230+
try {
1231+
((Session) session).insertRelationalTablets(Arrays.asList(first, second));
1232+
fail("Exception expected");
1233+
} catch (StatementExecutionException e) {
1234+
assertEquals("expected", e.getMessage());
1235+
}
1236+
1237+
assertEquals(0, (int) Whitebox.getInternalState(session, "mergeTabletsPerformanceCheckCount"));
1238+
assertEquals(0L, (long) Whitebox.getInternalState(session, "mergeTabletsCostInNanos"));
1239+
assertEquals(0L, (long) Whitebox.getInternalState(session, "insertTabletsCostInNanos"));
1240+
}
1241+
1242+
@Test
1243+
public void testInsertRelationalTabletsUseTableModelLeaderCache() throws Exception {
1244+
final Tablet tablet =
1245+
createRelationalTablet(
1246+
"table1",
1247+
Arrays.asList("tag1", "s1"),
1248+
Arrays.asList(TSDataType.STRING, TSDataType.INT64),
1249+
Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD),
1250+
1L,
1251+
"tag1",
1252+
11L);
1253+
final SessionConnection redirectedSessionConnection = Mockito.mock(SessionConnection.class);
1254+
final TEndPoint endPoint = new TEndPoint("127.0.0.2", 6667);
1255+
final Map<IDeviceID, TEndPoint> tableModelDeviceIdToEndpoint = new HashMap<>();
1256+
tableModelDeviceIdToEndpoint.put(tablet.getDeviceID(0), endPoint);
1257+
final Map<TEndPoint, SessionConnection> endPointToSessionConnection = new HashMap<>();
1258+
endPointToSessionConnection.put(endPoint, redirectedSessionConnection);
1259+
Whitebox.setInternalState(session, "enableRedirection", true);
1260+
Whitebox.setInternalState(
1261+
session, "tableModelDeviceIdToEndpoint", tableModelDeviceIdToEndpoint);
1262+
Whitebox.setInternalState(session, "endPointToSessionConnection", endPointToSessionConnection);
1263+
1264+
((Session) session).insertRelationalTablets(Collections.singletonList(tablet));
1265+
1266+
Mockito.verify(redirectedSessionConnection).insertTablets(any(TSInsertTabletsReq.class));
1267+
Mockito.verify(sessionConnection, Mockito.never()).insertTablets(any(TSInsertTabletsReq.class));
1268+
}
1269+
12021270
@Test
12031271
public void testTestInsertTablet() throws IoTDBConnectionException, StatementExecutionException {
12041272
List<IMeasurementSchema> schemas = new ArrayList<>();

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertTablets.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ public void validateTableSchema(Metadata metadata, MPPQueryContext context) {
7676
for (InsertTabletStatement insertTabletStatement :
7777
getInnerTreeStatement().getInsertTabletStatementList()) {
7878
final String database = AnalyzeUtils.getDatabaseName(insertTabletStatement, context);
79-
super.validateTableSchema(metadata, context, insertTabletStatement, database, false);
79+
super.validateTableSchema(metadata, context, insertTabletStatement, database, true);
8080
}
8181
}
8282

0 commit comments

Comments
 (0)