diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java index 5168678ff67d6..e305a13fdbd27 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java @@ -767,10 +767,30 @@ public void swapColumn(int src, int target) { if (inputLocations != null) { CommonUtils.swapArray(inputLocations, src, target); } + swapFailedMeasurementIndex(src, target); tagColumnIndices = null; attrColumnIndices = null; } + /** + * Keeps {@link #failedMeasurementIndex2Info} consistent with the measurement arrays after the + * columns at {@code src} and {@code target} are swapped, so that the failure records still point + * at the columns they were created for. + */ + private void swapFailedMeasurementIndex(final int src, final int target) { + if (failedMeasurementIndex2Info == null || failedMeasurementIndex2Info.isEmpty()) { + return; + } + final FailedMeasurementInfo srcInfo = failedMeasurementIndex2Info.remove(src); + final FailedMeasurementInfo targetInfo = failedMeasurementIndex2Info.remove(target); + if (targetInfo != null) { + failedMeasurementIndex2Info.put(src, targetInfo); + } + if (srcInfo != null) { + failedMeasurementIndex2Info.put(target, srcInfo); + } + } + /** * The oldToNewMapping array provides the mapping from old array positions to new array positions: * newArray[oldToNewMapping[oldIdx]] = oldArray[oldIdx] @@ -838,11 +858,37 @@ public void rebuildArraysAfterExpansion( typeConvertors = newTypeConvertors; inputLocations = newInputLocations; + rebuildFailedMeasurementIndex(newToOldMapping); + // Clear cached indices tagColumnIndices = null; attrColumnIndices = null; } + /** + * Rewrites the keys of {@link #failedMeasurementIndex2Info} from the old array positions to the + * new ones described by {@code newToOldMapping} ({@code newToOldMapping[newIdx] = oldIdx}), so + * that the failure records follow the reordered columns. + */ + private void rebuildFailedMeasurementIndex(final int[] newToOldMapping) { + if (failedMeasurementIndex2Info == null || failedMeasurementIndex2Info.isEmpty()) { + return; + } + final Map remapped = + new HashMap<>(failedMeasurementIndex2Info.size()); + for (int newIdx = 0; newIdx < newToOldMapping.length; newIdx++) { + final int oldIdx = newToOldMapping[newIdx]; + if (oldIdx < 0) { + continue; + } + final FailedMeasurementInfo info = failedMeasurementIndex2Info.get(oldIdx); + if (info != null) { + remapped.put(newIdx, info); + } + } + failedMeasurementIndex2Info = remapped; + } + private int getRequiredColumnArrayLength(final int index) { return Math.max(measurements == null ? 0 : measurements.length, index + 1); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertStatementPartialInsertTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertStatementPartialInsertTest.java index b6b26cacc8ecc..a37407636c978 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertStatementPartialInsertTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertStatementPartialInsertTest.java @@ -247,6 +247,34 @@ public void testRebuildArraysAfterExpansionHandlesShortValueArrays() throws Exce Assert.assertEquals(TsTableColumnCategory.TAG, statement.getColumnCategories()[0]); } + @Test + public void testSwapColumnRemapsFailedMeasurementIndex() throws Exception { + final InsertRowStatement statement = createInsertRowStatement(); + statement.markFailedMeasurement(0, new RuntimeException("failed")); + + statement.swapColumn(0, 1); + + Assert.assertEquals("s2", statement.getMeasurements()[0]); + Assert.assertNull(statement.getMeasurements()[1]); + Assert.assertEquals(1, statement.getFailedMeasurementInfoMap().size()); + Assert.assertEquals("s1", statement.getFailedMeasurementInfoMap().get(1).getMeasurement()); + } + + @Test + public void testRebuildArraysAfterExpansionRemapsFailedMeasurementIndex() throws Exception { + final InsertRowStatement statement = createInsertRowStatement(); + statement.markFailedMeasurement(0, new RuntimeException("failed")); + + // A failed column's name is null in the rebuilt measurement array, mirroring + // expandAndReorderTagColumns, which copies the (nulled) old measurement for the non-TAG area. + statement.rebuildArraysAfterExpansion(new int[] {-1, 0, 1}, new String[] {"tag1", null, "s2"}); + + Assert.assertEquals(1, statement.getFailedMeasurementInfoMap().size()); + Assert.assertEquals("s1", statement.getFailedMeasurementInfoMap().get(1).getMeasurement()); + Assert.assertNull(statement.getMeasurements()[1]); + Assert.assertEquals("s2", statement.getMeasurements()[2]); + } + @Test public void testInsertRowStatementSwapColumnInvalidatesTableDeviceId() throws Exception { final InsertRowStatement statement = createInsertRowStatement();