[Fix] Remap failed-measurement indices when reordering insert columns

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<Integer, FailedMeasurementInfo> 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);
   }

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();