This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch table_insert_tablets
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/table_insert_tablets by this
push:
new 8599e584630 fix review comments
8599e584630 is described below
commit 8599e58463038d6da8b6b9896de9f0ba0d2ed122
Author: Tian Jiang <[email protected]>
AuthorDate: Mon Jul 13 10:54:45 2026 +0800
fix review comments
---
.../java/org/apache/iotdb/session/Session.java | 223 +++++++++++++--------
.../apache/iotdb/session/SessionConnection.java | 6 +
.../java/org/apache/iotdb/session/SessionTest.java | 136 ++++++++++++-
.../plan/analyze/schema/SchemaValidator.java | 12 +-
.../write/RelationalInsertMultiTabletsNode.java | 4 +
.../plan/relational/planner/TableModelPlanner.java | 5 +-
.../planner/node/write/WritePlanNodeSplitTest.java | 43 ++++
7 files changed, 338 insertions(+), 91 deletions(-)
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 e39b5e8d377..63e11b9ba51 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
@@ -2924,14 +2924,20 @@ public class Session implements ISession {
updateRelationalTabletsReq(tabletGroup, connection, tablet);
return;
}
+ final List<SessionConnection> connections = new
ArrayList<>(tablet.getRowSize());
+ final Map<SessionConnection, Integer> rowCountByConnection = new
LinkedHashMap<>();
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));
+ final SessionConnection connection =
getTableModelSessionConnection(tablet.getDeviceID(row));
+ connections.add(connection);
+ rowCountByConnection.merge(connection, 1, Integer::sum);
+ }
+ final Map<SessionConnection, Tablet> subTabletByConnection = new
LinkedHashMap<>();
+ rowCountByConnection.forEach(
+ (connection, rowCount) ->
+ subTabletByConnection.put(
+ connection, createEmptyRelationalTabletLike(tablet,
rowCount)));
+ for (int row = 0; row < tablet.getRowSize(); row++) {
+ final Tablet subTablet = subTabletByConnection.get(connections.get(row));
for (int column = 0; column < subTablet.getSchemas().size(); column++) {
subTablet.addValue(
subTablet.getSchemas().get(column).getMeasurementName(),
@@ -2940,6 +2946,8 @@ public class Session implements ISession {
}
subTablet.addTimestamp(subTablet.getRowSize(), tablet.getTimestamp(row));
}
+ subTabletByConnection.forEach(
+ (connection, subTablet) -> updateRelationalTabletsReq(tabletGroup,
connection, subTablet));
}
private SessionConnection getTableModelSessionConnection(final IDeviceID
deviceID)
@@ -2950,6 +2958,10 @@ public class Session implements ISession {
}
private Tablet createEmptyRelationalTabletLike(final Tablet tablet) {
+ return createEmptyRelationalTabletLike(tablet, tablet.getRowSize());
+ }
+
+ private Tablet createEmptyRelationalTabletLike(final Tablet tablet, final
int rowSize) {
final List<String> measurements = new
ArrayList<>(tablet.getSchemas().size());
final List<TSDataType> dataTypes = new
ArrayList<>(tablet.getSchemas().size());
tablet
@@ -2960,11 +2972,7 @@ public class Session implements ISession {
dataTypes.add(schema.getType());
});
return new Tablet(
- tablet.getTableName(),
- measurements,
- dataTypes,
- tablet.getColumnTypes(),
- tablet.getRowSize());
+ tablet.getTableName(), measurements, dataTypes,
tablet.getColumnTypes(), rowSize);
}
private void updateRelationalTabletsReq(
@@ -2992,14 +3000,15 @@ public class Session implements ISession {
final SessionConnection connection = entry.getKey();
final RelationalTabletsReq relationalTabletsReq = entry.getValue();
try {
- connection.insertTablets(relationalTabletsReq.request);
+ connection.insertTablets(relationalTabletsReq.request,
relationalTabletsReq.deviceIDStrings);
} catch (RedirectException e) {
handleRelationalTabletsRedirection(e, relationalTabletsReq.deviceIDs);
} catch (IoTDBConnectionException e) {
if (endPointToSessionConnection != null &&
endPointToSessionConnection.size() > 1) {
removeBrokenSessionConnection(connection);
try {
-
getDefaultSessionConnection().insertTablets(relationalTabletsReq.request);
+ getDefaultSessionConnection()
+ .insertTablets(relationalTabletsReq.request,
relationalTabletsReq.deviceIDStrings);
} catch (RedirectException ignored) {
}
} else {
@@ -3024,7 +3033,8 @@ public class Session implements ISession {
return CompletableFuture.runAsync(
() -> {
try {
-
connection.insertTablets(relationalTabletsReq.request);
+ connection.insertTablets(
+ relationalTabletsReq.request,
relationalTabletsReq.deviceIDStrings);
} catch (RedirectException e) {
handleRelationalTabletsRedirection(e,
relationalTabletsReq.deviceIDs);
} catch (StatementExecutionException e) {
@@ -3033,7 +3043,9 @@ public class Session implements ISession {
removeBrokenSessionConnection(connection);
try {
getDefaultSessionConnection()
- .insertTablets(relationalTabletsReq.request);
+ .insertTablets(
+ relationalTabletsReq.request,
+ relationalTabletsReq.deviceIDStrings);
} catch (IoTDBConnectionException |
StatementExecutionException ex) {
throw new CompletionException(ex);
} catch (RedirectException ignored) {
@@ -3067,6 +3079,16 @@ public class Session implements ISession {
private void handleRelationalTabletsRedirection(
final RedirectException redirectException, final List<IDeviceID>
deviceIDs) {
+ final Map<String, TEndPoint> deviceEndPointMap =
redirectException.getDeviceEndPointMap();
+ if (deviceEndPointMap != null) {
+ for (final IDeviceID deviceID : deviceIDs) {
+ final TEndPoint endPoint = deviceEndPointMap.get(deviceID.toString());
+ if (endPoint != null) {
+ handleRedirection(deviceID, endPoint);
+ }
+ }
+ return;
+ }
final List<TEndPoint> endPointList = redirectException.getEndPointList();
if (endPointList == null) {
return;
@@ -3087,23 +3109,52 @@ public class Session implements ISession {
sortedTablets.sort(
Comparator.comparing(
Tablet::getTableName,
Comparator.nullsFirst(Comparator.naturalOrder())));
- final List<RelationalTabletWithColumnIndexMap> mergedTablets = new
ArrayList<>(tablets.size());
+ final List<Tablet> mergedTablets = new ArrayList<>(tablets.size());
+ final List<Tablet> currentGroup = new ArrayList<>();
+ Map<String, Integer> currentColumnIndexMap = new LinkedHashMap<>();
+ List<String> currentMeasurements = new ArrayList<>();
+ List<TSDataType> currentDataTypes = new ArrayList<>();
+ List<ColumnCategory> currentColumnTypes = new ArrayList<>();
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 (currentGroup.isEmpty()) {
+ currentGroup.add(tablet);
+ currentColumnIndexMap = getColumnIndexMap(tablet);
+ currentMeasurements = getMeasurements(tablet);
+ currentDataTypes = getDataTypes(tablet);
+ currentColumnTypes = new ArrayList<>(tablet.getColumnTypes());
+ continue;
}
- if (!merged) {
- mergedTablets.add(
- new RelationalTabletWithColumnIndexMap(tablet,
getColumnIndexMap(tablet)));
+ if (canMergeRelationalTablets(
+ currentGroup.get(0).getTableName(),
+ currentColumnIndexMap,
+ currentDataTypes,
+ currentColumnTypes,
+ tablet)) {
+ currentGroup.add(tablet);
+ addMissingColumns(
+ tablet,
+ currentColumnIndexMap,
+ currentMeasurements,
+ currentDataTypes,
+ currentColumnTypes);
+ anyMerged = true;
+ continue;
}
+ mergedTablets.add(
+ mergeRelationalTablets(
+ currentGroup, currentMeasurements, currentDataTypes,
currentColumnTypes));
+ currentGroup.clear();
+ currentGroup.add(tablet);
+ currentColumnIndexMap = getColumnIndexMap(tablet);
+ currentMeasurements = getMeasurements(tablet);
+ currentDataTypes = getDataTypes(tablet);
+ currentColumnTypes = new ArrayList<>(tablet.getColumnTypes());
+ }
+ if (!currentGroup.isEmpty()) {
+ mergedTablets.add(
+ mergeRelationalTablets(
+ currentGroup, currentMeasurements, currentDataTypes,
currentColumnTypes));
}
if (anyMerged) {
consecutiveRelationalTabletMergeMissCount = 0;
@@ -3114,9 +3165,7 @@ public class Session implements ISession {
enableMergeTablets = false;
}
}
- return mergedTablets.stream()
- .map(relationalTablet -> relationalTablet.tablet)
- .collect(Collectors.toList());
+ return mergedTablets;
}
private void recordMergeTabletsCost(
@@ -3137,8 +3186,12 @@ public class Session implements ISession {
}
private boolean canMergeRelationalTablets(
- final Tablet left, final Map<String, Integer> leftColumnIndexMap, final
Tablet right) {
- if (!Objects.equals(left.getTableName(), right.getTableName())) {
+ final String leftTableName,
+ final Map<String, Integer> leftColumnIndexMap,
+ final List<TSDataType> leftDataTypes,
+ final List<ColumnCategory> leftColumnTypes,
+ final Tablet right) {
+ if (!Objects.equals(leftTableName, right.getTableName())) {
return false;
}
final List<IMeasurementSchema> rightSchemas = right.getSchemas();
@@ -3149,13 +3202,13 @@ public class Session implements ISession {
if (leftIndex == null) {
continue;
}
- if (left.getSchemas().get(leftIndex).getType() != rightSchema.getType()
- || left.getColumnTypes().get(leftIndex) !=
right.getColumnTypes().get(rightIndex)) {
+ if (leftDataTypes.get(leftIndex) != rightSchema.getType()
+ || leftColumnTypes.get(leftIndex) !=
right.getColumnTypes().get(rightIndex)) {
return false;
}
duplicatedColumnCount++;
}
- return (double) duplicatedColumnCount / Math.max(left.getSchemas().size(),
rightSchemas.size())
+ return (double) duplicatedColumnCount / Math.max(leftDataTypes.size(),
rightSchemas.size())
> 0.5;
}
@@ -3167,53 +3220,66 @@ public class Session implements ISession {
return columnIndexMap;
}
- private Map<String, Integer> getColumnIndexMap(final List<String>
measurements) {
- final Map<String, Integer> columnIndexMap = new
HashMap<>(measurements.size());
- for (int i = 0; i < measurements.size(); i++) {
- columnIndexMap.put(measurements.get(i), i);
- }
- return columnIndexMap;
+ private List<String> getMeasurements(final Tablet tablet) {
+ final List<String> measurements = new
ArrayList<>(tablet.getSchemas().size());
+ tablet.getSchemas().forEach(schema ->
measurements.add(schema.getMeasurementName()));
+ return measurements;
}
- private RelationalTabletWithColumnIndexMap mergeRelationalTablets(
- final RelationalTabletWithColumnIndexMap left, final Tablet right) {
- final List<String> measurements =
- new ArrayList<>(left.tablet.getSchemas().size() +
right.getSchemas().size());
- final List<TSDataType> dataTypes =
- new ArrayList<>(left.tablet.getSchemas().size() +
right.getSchemas().size());
- final List<ColumnCategory> 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));
+ private List<TSDataType> getDataTypes(final Tablet tablet) {
+ final List<TSDataType> dataTypes = new
ArrayList<>(tablet.getSchemas().size());
+ tablet.getSchemas().forEach(schema -> dataTypes.add(schema.getType()));
+ return dataTypes;
+ }
+
+ private void addMissingColumns(
+ final Tablet tablet,
+ final Map<String, Integer> columnIndexMap,
+ final List<String> measurements,
+ final List<TSDataType> dataTypes,
+ final List<ColumnCategory> columnTypes) {
+ for (int i = 0; i < tablet.getSchemas().size(); i++) {
+ final IMeasurementSchema schema = tablet.getSchemas().get(i);
+ if (columnIndexMap.containsKey(schema.getMeasurementName())) {
+ continue;
}
+ columnIndexMap.put(schema.getMeasurementName(), measurements.size());
+ measurements.add(schema.getMeasurementName());
+ dataTypes.add(schema.getType());
+ columnTypes.add(tablet.getColumnTypes().get(i));
}
+ }
- final int rowCount = left.tablet.getRowSize() + right.getRowSize();
+ private Tablet mergeRelationalTablets(
+ final List<Tablet> tablets,
+ final List<String> measurements,
+ final List<TSDataType> dataTypes,
+ final List<ColumnCategory> columnTypes) {
+ if (tablets.size() == 1) {
+ return tablets.get(0);
+ }
+ int rowCount = 0;
+ for (final Tablet tablet : tablets) {
+ rowCount += tablet.getRowSize();
+ }
final Tablet merged =
- new Tablet(left.tablet.getTableName(), measurements, dataTypes,
columnTypes, rowCount);
+ new Tablet(tablets.get(0).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);
+ int targetRowOffset = 0;
+ for (final Tablet tablet : tablets) {
+ copyTabletColumns(tablet, targetRowOffset, measurements, values,
bitMaps, timestamps);
+ targetRowOffset += tablet.getRowSize();
+ }
merged.setTimestamps(timestamps);
merged.setValues(values);
merged.setBitMaps(bitMaps);
merged.setRowSize(rowCount);
- return new RelationalTabletWithColumnIndexMap(merged,
getColumnIndexMap(measurements));
+ return merged;
}
private void copyTabletColumns(
@@ -3295,24 +3361,12 @@ public class Session implements ISession {
}
}
- private static class RelationalTabletWithColumnIndexMap {
-
- private final Tablet tablet;
- private final Map<String, Integer> columnIndexMap;
-
- private RelationalTabletWithColumnIndexMap(
- final Tablet tablet, final Map<String, Integer> columnIndexMap) {
- this.tablet = tablet;
- this.columnIndexMap = columnIndexMap;
- }
- }
-
private class RelationalTabletsReq {
private final TSInsertTabletsReq request = new TSInsertTabletsReq();
private final List<IDeviceID> deviceIDs = new ArrayList<>();
+ private final List<String> deviceIDStrings = new ArrayList<>();
private final List<Tablet> tablets = new ArrayList<>();
- private final Map<IDeviceID, Tablet> tabletByDevice = new
LinkedHashMap<>();
private void addTablet(final Tablet tablet) {
tablets.add(tablet);
@@ -3321,13 +3375,16 @@ public class Session implements ISession {
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));
+ for (int row = 0; row < tablet.getRowSize(); row++) {
+ final IDeviceID deviceID = tablet.getDeviceID(row);
+ deviceIDs.add(deviceID);
+ deviceIDStrings.add(deviceID.toString());
+ }
}
}
diff --git
a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
index ad6699aaf87..4c30ba468f8 100644
---
a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
+++
b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
@@ -870,6 +870,12 @@ public class SessionConnection {
() -> insertTabletsInternal(request), request::getPrefixPaths);
}
+ protected void insertTablets(TSInsertTabletsReq request, List<String>
devices)
+ throws IoTDBConnectionException, StatementExecutionException,
RedirectException {
+ callWithRetryAndVerifyWithRedirectionForMultipleDevices(
+ () -> insertTabletsInternal(request), () -> devices);
+ }
+
private TSStatus insertTabletsInternal(TSInsertTabletsReq request) throws
TException {
request.setSessionId(sessionId);
return client.insertTablets(request);
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 0b317311ce7..ac5959bd3e9 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
@@ -25,6 +25,7 @@ import org.apache.iotdb.isession.ISession;
import org.apache.iotdb.isession.SessionConfig;
import org.apache.iotdb.isession.util.Version;
import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.RedirectException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq;
import org.apache.iotdb.service.rpc.thrift.TSQueryTemplateResp;
@@ -61,6 +62,7 @@ import java.util.Map;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
public class SessionTest {
@@ -1029,6 +1031,72 @@ public class SessionTest {
Assert.assertTrue((boolean) mergedTablet.getValue(1, 3));
}
+ @Test
+ public void testMergeRelationalTabletsChecksColumnsAddedByEarlierTablets()
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 Tablet conflicting =
+ createRelationalTablet(
+ "table1",
+ Arrays.asList("tag1", "s1", "s3"),
+ Arrays.asList(TSDataType.STRING, TSDataType.INT64,
TSDataType.INT32),
+ Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD,
ColumnCategory.FIELD),
+ 3,
+ "d3",
+ 33L,
+ 3);
+
+ final List<Tablet> mergedTablets =
+ Whitebox.invokeMethod(
+ session, "mergeRelationalTablets", Arrays.asList(first, second,
conflicting));
+
+ assertEquals(2, mergedTablets.size());
+ assertEquals(2, mergedTablets.get(0).getRowSize());
+ Assert.assertSame(conflicting, mergedTablets.get(1));
+ }
+
+ @Test(timeout = 10000)
+ public void testMergeLargeRelationalTabletList() throws Exception {
+ final int tabletCount = 5000;
+ final List<Tablet> tablets = new ArrayList<>(tabletCount);
+ for (int i = 0; i < tabletCount; i++) {
+ tablets.add(
+ createRelationalTablet(
+ "table1",
+ Arrays.asList("tag1", "s1"),
+ Arrays.asList(TSDataType.STRING, TSDataType.INT64),
+ Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD),
+ i,
+ "d" + i,
+ (long) i));
+ }
+
+ final List<Tablet> mergedTablets =
+ Whitebox.invokeMethod(session, "mergeRelationalTablets", tablets);
+
+ assertEquals(1, mergedTablets.size());
+ assertEquals(tabletCount, mergedTablets.get(0).getRowSize());
+ assertEquals((long) tabletCount - 1,
mergedTablets.get(0).getValue(tabletCount - 1, 1));
+ }
+
@Test
public void
testMergeRelationalTabletsSkipLowDuplicatedColumnsAndDifferentTables()
throws Exception {
@@ -1205,7 +1273,7 @@ public class SessionTest {
public void testMergeTabletsCostNotRecordedWhenInsertFails() throws
Exception {
Mockito.doThrow(new StatementExecutionException("expected"))
.when(sessionConnection)
- .insertTablets(any(TSInsertTabletsReq.class));
+ .insertTablets(any(TSInsertTabletsReq.class), anyList());
final Tablet first =
createRelationalTablet(
"table1",
@@ -1263,8 +1331,70 @@ public class SessionTest {
((Session)
session).insertRelationalTablets(Collections.singletonList(tablet));
-
Mockito.verify(redirectedSessionConnection).insertTablets(any(TSInsertTabletsReq.class));
- Mockito.verify(sessionConnection,
Mockito.never()).insertTablets(any(TSInsertTabletsReq.class));
+ Mockito.verify(redirectedSessionConnection)
+ .insertTablets(any(TSInsertTabletsReq.class), anyList());
+ Mockito.verify(sessionConnection, Mockito.never())
+ .insertTablets(any(TSInsertTabletsReq.class), anyList());
+ }
+
+ @Test
+ public void testInsertRelationalTabletsUpdatesTableModelLeaderCache() 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 IDeviceID deviceID = tablet.getDeviceID(0);
+ final TEndPoint redirectEndPoint = new TEndPoint("127.0.0.2", 6667);
+ final Map<TEndPoint, SessionConnection> endPointToSessionConnection = new
HashMap<>();
+ endPointToSessionConnection.put(redirectEndPoint,
Mockito.mock(SessionConnection.class));
+ Whitebox.setInternalState(session, "enableRedirection", true);
+ Whitebox.setInternalState(session, "tableModelDeviceIdToEndpoint", new
HashMap<>());
+ Whitebox.setInternalState(session, "endPointToSessionConnection",
endPointToSessionConnection);
+ Mockito.doThrow(
+ new
RedirectException(Collections.singletonMap(deviceID.toString(),
redirectEndPoint)))
+ .when(sessionConnection)
+ .insertTablets(any(TSInsertTabletsReq.class), anyList());
+
+ ((Session)
session).insertRelationalTablets(Collections.singletonList(tablet));
+
+ final Map<IDeviceID, TEndPoint> tableModelDeviceIdToEndpoint =
+ Whitebox.getInternalState(session, "tableModelDeviceIdToEndpoint");
+ assertEquals(redirectEndPoint, tableModelDeviceIdToEndpoint.get(deviceID));
+
Mockito.verify(sessionConnection).insertTablets(any(TSInsertTabletsReq.class),
anyList());
+ }
+
+ @Test
+ public void testSplitRelationalTabletAllocatesOncePerEndpoint() throws
Exception {
+ final int rowCount = 1000;
+ final List<String> measurements = Arrays.asList("tag1", "s1");
+ final Tablet tablet =
+ new Tablet(
+ "table1",
+ measurements,
+ Arrays.asList(TSDataType.STRING, TSDataType.INT64),
+ Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD),
+ rowCount);
+ for (int row = 0; row < rowCount; row++) {
+ tablet.addTimestamp(row, row);
+ tablet.addValue(measurements.get(0), row, "d" + row);
+ tablet.addValue(measurements.get(1), row, (long) row);
+ }
+ Whitebox.setInternalState(session, "tableModelDeviceIdToEndpoint", new
HashMap<>());
+ final Map<SessionConnection, Object> tabletGroup = new HashMap<>();
+
+ Whitebox.invokeMethod(session, "addRelationalTabletToGroup", tabletGroup,
tablet);
+
+ assertEquals(1, tabletGroup.size());
+ final List<Tablet> groupedTablets =
+ Whitebox.getInternalState(tabletGroup.values().iterator().next(),
"tablets");
+ assertEquals(1, groupedTablets.size());
+ assertEquals(rowCount, groupedTablets.get(0).getMaxRowNumber());
+ assertEquals(rowCount, groupedTablets.get(0).getRowSize());
}
@Test
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 f400ec4f3cd..b80d716a1ac 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
@@ -41,7 +41,9 @@ import org.apache.tsfile.file.metadata.enums.TSEncoding;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.HashSet;
import java.util.List;
+import java.util.Set;
import static org.apache.iotdb.commons.utils.PathUtils.unQualifyDatabaseName;
@@ -93,13 +95,17 @@ public class SchemaValidator {
final MPPQueryContext context,
final AccessControl accessControl) {
try {
+ final String database =
unQualifyDatabaseName(insertStatement.getDatabase());
+ final Set<String> checkedTables = new HashSet<>();
for (InsertTabletStatement tabletStatement :
insertStatement.getInnerTreeStatement().getInsertTabletStatementList()) {
+ final String tableName = tabletStatement.getDevicePath().getFullPath();
+ if (!checkedTables.add(tableName)) {
+ continue;
+ }
accessControl.checkCanInsertIntoTable(
context.getSession().getUserName(),
- new QualifiedObjectName(
- unQualifyDatabaseName(insertStatement.getDatabase()),
- tabletStatement.getDevicePath().getFullPath()),
+ new QualifiedObjectName(database, tableName),
context);
}
insertStatement.validateTableSchema(metadata, context);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertMultiTabletsNode.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertMultiTabletsNode.java
index 7ce55c4cc87..261912b8edd 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertMultiTabletsNode.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertMultiTabletsNode.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write;
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
import
org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
@@ -58,9 +59,11 @@ public class RelationalInsertMultiTabletsNode extends
InsertMultiTabletsNode {
@Override
public List<WritePlanNode> splitByPartition(IAnalysis analysis) {
Map<TRegionReplicaSet, RelationalInsertMultiTabletsNode> splitMap = new
HashMap<>();
+ final List<TEndPoint> redirectNodeList = new ArrayList<>();
for (int i = 0; i < insertTabletNodeList.size(); i++) {
InsertTabletNode insertTabletNode = insertTabletNodeList.get(i);
List<WritePlanNode> tmpResult =
insertTabletNode.splitByPartition(analysis);
+ redirectNodeList.addAll(analysis.getRedirectNodeList());
for (WritePlanNode subNode : tmpResult) {
TRegionReplicaSet dataRegionReplicaSet = ((InsertNode)
subNode).getDataRegionReplicaSet();
RelationalInsertMultiTabletsNode tmpNode =
splitMap.get(dataRegionReplicaSet);
@@ -75,6 +78,7 @@ public class RelationalInsertMultiTabletsNode extends
InsertMultiTabletsNode {
tmpNode.addInsertTabletNode((InsertTabletNode) subNode, i);
}
}
+ analysis.setRedirectNodeList(redirectNodeList);
return new ArrayList<>(splitMap.values());
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java
index c8433d95416..9a32bb52860 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java
@@ -55,6 +55,7 @@ import
org.apache.iotdb.db.queryengine.plan.scheduler.ClusterScheduler;
import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler;
import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler;
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.InsertTabletStatement;
import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
@@ -238,8 +239,8 @@ public class TableModelPlanner implements IPlanner {
((WrappedInsertStatement) statementToRedirect).getInnerTreeStatement();
if (!analysis.isFinishQueryAfterAnalyze()) {
- // Table Model Session only supports insertTablet
- if (insertStatement instanceof InsertTabletStatement) {
+ if (insertStatement instanceof InsertTabletStatement
+ || insertStatement instanceof InsertMultiTabletsStatement) {
if (tsstatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode())
{
boolean needRedirect = false;
List<TEndPoint> redirectNodeList = analysis.getRedirectNodeList();
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/write/WritePlanNodeSplitTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/write/WritePlanNodeSplitTest.java
index 6d5185af4ba..1805b4e128a 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/write/WritePlanNodeSplitTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/write/WritePlanNodeSplitTest.java
@@ -43,6 +43,7 @@ import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNod
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsNode;
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsOfOneDeviceNode;
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode;
+import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertMultiTabletsNode;
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode;
import org.apache.tsfile.enums.TSDataType;
@@ -364,6 +365,48 @@ public class WritePlanNodeSplitTest {
}
}
+ @Test
+ public void testSplitRelationalInsertMultiTabletsCollectsRedirectNodes()
+ throws IllegalPathException {
+ final RelationalInsertMultiTabletsNode insertMultiTabletsNode =
+ new RelationalInsertMultiTabletsNode(new PlanNodeId("plan node 2"));
+ final List<DataPartitionQueryParam> dataPartitionQueryParams = new
ArrayList<>();
+ for (int tabletIndex = 0; tabletIndex < 2; tabletIndex++) {
+ final RelationalInsertTabletNode insertTabletNode =
+ new RelationalInsertTabletNode(new PlanNodeId("plan node 2"));
+ insertTabletNode.setTargetPath(new PartialPath("root.sg1"));
+ insertTabletNode.setTimes(new long[] {1, 101});
+ insertTabletNode.setDataTypes(new TSDataType[] {TSDataType.STRING,
TSDataType.INT32});
+ insertTabletNode.setColumns(
+ new Object[] {
+ new Binary[] {
+ new Binary("d" + tabletIndex + "0", StandardCharsets.UTF_8),
+ new Binary("d" + tabletIndex + "1", StandardCharsets.UTF_8)
+ },
+ new int[] {tabletIndex * 2, tabletIndex * 2 + 1}
+ });
+ insertTabletNode.setColumnCategories(
+ new TsTableColumnCategory[] {TsTableColumnCategory.TAG,
TsTableColumnCategory.FIELD});
+ insertTabletNode.setRowCount(2);
+ insertMultiTabletsNode.addInsertTabletNode(insertTabletNode,
tabletIndex);
+ for (int row = 0; row < insertTabletNode.getRowCount(); row++) {
+ final DataPartitionQueryParam queryParam = new
DataPartitionQueryParam();
+ queryParam.setDeviceID(insertTabletNode.getDeviceID(row));
+
queryParam.setTimePartitionSlotList(insertTabletNode.getTimePartitionSlots());
+ dataPartitionQueryParams.add(queryParam);
+ }
+ }
+ final Analysis analysis = new Analysis();
+ analysis.setDataPartitionInfo(getDataPartition(dataPartitionQueryParams));
+
+ insertMultiTabletsNode.splitByPartition(analysis);
+
+ Assert.assertEquals(4, analysis.getRedirectNodeList().size());
+ for (final TEndPoint redirectNode : analysis.getRedirectNodeList()) {
+ Assert.assertNotNull(redirectNode);
+ }
+ }
+
@Test
public void testInsertMultiTablets() throws IllegalPathException {
InsertMultiTabletsNode insertMultiTabletsNode =
