Fix Prepared Table-Model INSERT VALUES Execution

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBSetSystemStatusTableIT.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
 import org.apache.iotdb.itbase.env.BaseEnv;
 import org.apache.iotdb.itbase.exception.InconsistentDataException;
+import org.apache.iotdb.itbase.runtime.ClusterTestConnection;
 
 import org.awaitility.Awaitility;
 import org.junit.AfterClass;
@@ -52,6 +53,40 @@ public static void tearDown() throws Exception {
     EnvFactory.getEnv().cleanClusterEnvironment();
   }
 
+  @Test
+  public void testPreparedInsertValuesExecution() throws Exception {
+    try (Connection tableConnection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement table = tableConnection.createStatement();
+        Statement prepared =
+            (tableConnection instanceof ClusterTestConnection
+                    ? ((ClusterTestConnection) tableConnection)
+                        .writeConnection.getUnderlyingConnection()
+                    : tableConnection)
+                .createStatement()) {
+      table.execute("create database prepared_insert_execution");
+      table.execute("use prepared_insert_execution");
+      table.execute("create table prepared_t(d STRING TAG,a STRING ATTRIBUTE,s INT32 FIELD)");
+      prepared.execute(
+          "prepare prepared_insert from insert into prepared_t(time,d,a,s) "
+              + "values(1,'p1','prepared',11),(2,'p2','prepared',12)");
+      prepared.execute("execute prepared_insert");
+      table.execute("update prepared_t set a='changed' where d='p1'");
+      prepared.execute("execute prepared_insert");
+      prepared.execute(
+          "execute immediate 'insert into prepared_t(time,d,a,s) "
+              + "values(3,''p3'',''immediate'',13)'");
+      try (ResultSet result = table.executeQuery("select d,a,s from prepared_t order by time")) {
+        for (int i = 1; i <= 3; i++) {
+          Assert.assertTrue(result.next());
+          Assert.assertEquals("p" + i, result.getString(1));
+          Assert.assertEquals(i == 3 ? "immediate" : "prepared", result.getString(2));
+          Assert.assertEquals(10 + i, result.getInt(3));
+        }
+        Assert.assertFalse(result.next());
+      }
+    }
+  }
+
   @Test
   public void setSystemStatus() {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java

@@ -103,6 +103,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ExplainOutputFormat;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ExtendRegion;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Flush;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRows;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.KillQuery;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadConfiguration;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadModel;
@@ -815,6 +816,10 @@ private IQueryExecution createQueryExecutionForTableModel(
 
       org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement resolvedSql =
           preparedInfo.getSql();
+      if (resolvedSql instanceof InsertRows) {
+        resolvedSql = ((InsertRows) resolvedSql).copyForPreparedExecution();
+      }
+      innerStatement = resolvedSql;
       parameterLookup =
           ParameterExtractor.bindParameters(resolvedSql, executeStatement.getParameters());
       parameters = new ArrayList<>(executeStatement.getParameters());
@@ -838,6 +843,7 @@ private IQueryExecution createQueryExecutionForTableModel(
 
       org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement resolvedSql =
           sqlParser.createStatement(sql, clientSession.getZoneId(), clientSession);
+      innerStatement = resolvedSql;
 
       if (!literalParameters.isEmpty()) {
         parameterLookup = ParameterExtractor.bindParameters(resolvedSql, literalParameters);
@@ -856,8 +862,8 @@ private IQueryExecution createQueryExecutionForTableModel(
       }
     }
 
-    if (statement instanceof WrappedInsertStatement) {
-      ((WrappedInsertStatement) statement).setContext(queryContext);
+    if (innerStatement instanceof WrappedInsertStatement) {
+      ((WrappedInsertStatement) innerStatement).setContext(queryContext);
     }
 
     // Create QueryExecution with TableModelPlanner

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

@@ -20,6 +20,7 @@
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
 import org.apache.iotdb.calc.exception.QueryProcessException;
+import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IAstVisitor;
 import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
@@ -32,6 +33,7 @@
 
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.utils.Constants;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -44,6 +46,9 @@
 
 public class InsertRows extends WrappedInsertStatement {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(InsertRows.class);
+
   // Only InsertRows constructed by Pipe will be set to true
   private boolean allowCreateTable = false;
 
@@ -61,6 +66,46 @@ public InsertRowsStatement getInnerTreeStatement() {
     return ((InsertRowsStatement) super.getInnerTreeStatement());
   }
 
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE + getInnerTreeStatement().ramBytesUsed();
+  }
+
+  /** Copy the parsed VALUES payload so analysis cannot mutate the prepared statement. */
+  public InsertRows copyForPreparedExecution() {
+    InsertRowsStatement original = getInnerTreeStatement();
+    List<InsertRowStatement> rows = new ArrayList<>(original.getInsertRowStatementList().size());
+    for (InsertRowStatement row : original.getInsertRowStatementList()) {
+      InsertRowStatement copy = new InsertRowStatement();
+      copy.setDevicePath(new PartialPath(row.getDevicePath().getNodes().clone()));
+      copy.setTime(row.getTime());
+      copy.setMeasurements(copyArray(row.getMeasurements()));
+      copy.setMeasurementSchemas(copyArray(row.getMeasurementSchemas()));
+      copy.setDataTypes(copyArray(row.getDataTypes()));
+      copy.setColumnCategories(copyArray(row.getColumnCategories()));
+      copy.setValues(copyArray(row.getValues()));
+      copy.setNeedInferType(row.isNeedInferType());
+      copy.setAligned(row.isAligned());
+      copy.setWriteToTable(row.isWriteToTable());
+      if (row.getFailedMeasurementInfoMap() != null) {
+        copy.setFailedMeasurementIndex2Info(new HashMap<>(row.getFailedMeasurementInfoMap()));
+      }
+      row.getDatabaseName().ifPresent(copy::setDatabaseName);
+      rows.add(copy);
+    }
+    InsertRowsStatement copy = new InsertRowsStatement();
+    copy.setInsertRowStatementList(rows);
+    copy.setUniformTargets(original.isUniformTargets());
+    copy.setWriteToTable(original.isWriteToTable());
+    InsertRows result = new InsertRows(copy, null);
+    result.setAllowCreateTable(allowCreateTable);
+    return result;
+  }
+
+  private static <T> T[] copyArray(T[] array) {
+    return array == null ? null : array.clone();
+  }
+
   @Override
   public void updateAfterSchemaValidation(MPPQueryContext context) throws QueryProcessException {
     getInnerTreeStatement().updateAfterSchemaValidation(context);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java

@@ -53,6 +53,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -579,7 +580,10 @@ public void removeAttributeColumns() {
 
   protected abstract void subRemoveAttributeColumns(List<Integer> columnsToKeep);
 
-  public static class FailedMeasurementInfo {
+  public static class FailedMeasurementInfo implements Accountable {
+    private static final long INSTANCE_SIZE =
+        RamUsageEstimator.shallowSizeOfInstance(FailedMeasurementInfo.class);
+
     protected String measurement;
     protected TSDataType dataType;
     protected Object value;
@@ -604,6 +608,21 @@ public TSDataType getDataType() {
     public Object getValue() {
       return value;
     }
+
+    @Override
+    public long ramBytesUsed() {
+      long size =
+          INSTANCE_SIZE
+              + RamUsageEstimator.sizeOf(measurement)
+              + RamUsageEstimator.sizeOfObject(value);
+      // Exceptions can retain additional copies of the failed input in their messages.
+      Set<Throwable> visited = Collections.newSetFromMap(new IdentityHashMap<>());
+      for (Throwable error = cause; error != null && visited.add(error); error = error.getCause()) {
+        size +=
+            RamUsageEstimator.sizeOfObject(error) + RamUsageEstimator.sizeOf(error.getMessage());
+      }
+      return size;
+    }
   }
 
   // endregion
@@ -1015,6 +1034,7 @@ public long ramBytesUsed() {
             + RamUsageEstimator.sizeOf(measurements)
             + RamUsageEstimator.shallowSizeOf(dataTypes)
             + RamUsageEstimator.shallowSizeOf(columnCategories)
+            + RamUsageEstimator.sizeOfMap(failedMeasurementIndex2Info)
             // We assume that the integers are all cached by JVM
             + shallowSizeOfList(tagColumnIndices)
             + shallowSizeOfList(attrColumnIndices)

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertRowsStatement.java

@@ -233,8 +233,12 @@ protected long calculateBytesUsed() {
     return INSTANCE_SIZE
         + (Objects.nonNull(insertRowStatementList)
             ? UpdateDetailContainer.LIST_SIZE
+                + RamUsageEstimator.alignObjectSize(
+                    RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+                        + (long) RamUsageEstimator.NUM_BYTES_OBJECT_REF
+                            * insertRowStatementList.size())
                 + insertRowStatementList.stream()
-                    .mapToLong(InsertRowStatement::calculateBytesUsed)
+                    .mapToLong(InsertRowStatement::ramBytesUsed)
                     .reduce(0L, Long::sum)
             : 0);
   }