diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBSetSystemStatusTableIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBSetSystemStatusTableIT.java index 4dac5e9feb4ec..685610a663592 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBSetSystemStatusTableIT.java +++ b/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); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java index 79e6e38d4fe1f..eb3d0e6869871 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java +++ b/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 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertRows.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertRows.java index 6b5804ba4ebb8..f40c8ca0350be 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/InsertRows.java +++ b/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 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[] copyArray(T[] array) { + return array == null ? null : array.clone(); + } + @Override public void updateAfterSchemaValidation(MPPQueryContext context) throws QueryProcessException { getInnerTreeStatement().updateAfterSchemaValidation(context); 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 e305a13fdbd27..385605e6cad96 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 @@ -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 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 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) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertRowsStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertRowsStatement.java index 728831a7d443b..e89d32231c536 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertRowsStatement.java +++ b/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); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/sql/AstMemoryEstimatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/sql/AstMemoryEstimatorTest.java index 0738b8cd27acc..c107593c84feb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/sql/AstMemoryEstimatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/sql/AstMemoryEstimatorTest.java @@ -19,17 +19,37 @@ package org.apache.iotdb.db.queryengine.plan.relational.sql; +import org.apache.iotdb.commons.memory.IMemoryBlock; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement; +import org.apache.iotdb.commons.schema.table.TsTable; +import org.apache.iotdb.commons.schema.table.column.FieldColumnSchema; +import org.apache.iotdb.commons.schema.table.column.TagColumnSchema; +import org.apache.iotdb.commons.schema.table.column.TimeColumnSchema; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.protocol.session.ClientSession; import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.protocol.session.InternalClientSession; +import org.apache.iotdb.db.protocol.session.PreparedStatementInfo; +import org.apache.iotdb.db.queryengine.plan.Coordinator; +import org.apache.iotdb.db.queryengine.plan.execution.config.session.PreparedStatementHelper; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.InsertRows; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Prepare; import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; +import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache; +import org.apache.iotdb.db.schemaengine.table.ITableCache; +import org.apache.tsfile.enums.TSDataType; import org.junit.Before; import org.junit.Test; import java.time.ZoneId; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; public class AstMemoryEstimatorTest { @@ -756,6 +776,154 @@ public void testPrepare() { assertTrue("Memory size should be positive", memorySize > 0); } + @Test + public void testPreparedInsertValuesMemoryAccounting() { + TsTable table = new TsTable("insert_memory"); + table.addColumnSchema(new TimeColumnSchema("time", TSDataType.TIMESTAMP)); + table.addColumnSchema(new TagColumnSchema("device", TSDataType.STRING)); + table.addColumnSchema(new FieldColumnSchema("value", TSDataType.STRING)); + ITableCache cache = DataNodeTableCache.getInstance(); + cache.preUpdateTable("testdb", table, null); + cache.commitUpdateTable("testdb", "insert_memory", null, null); + ClientSession session = new ClientSession(null); + IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + int originalDataNodeId = config.getDataNodeId(); + config.setDataNodeId(0); + try { + IMemoryBlock memory = Coordinator.getCoordinatorMemoryBlock(); + long initialUsage = memory.getUsedMemoryInBytes(); + Prepare prepare = + (Prepare) + parseSQL( + "PREPARE insert_values FROM INSERT INTO insert_memory(time,device,value) " + + "VALUES(1,'d1','small'),(2,'d2',null)"); + InsertRows insert = (InsertRows) prepare.getSql(); + long size = insert.ramBytesUsed(); + long rowSizes = + insert.getInnerTreeStatement().getInsertRowStatementList().stream() + .mapToLong(row -> row.ramBytesUsed()) + .sum(); + assertTrue(size > rowSizes); + assertTrue(prepare.ramBytesUsed() > size); + assertEquals(size, insert.ramBytesUsed()); + + Statement singleRow = + parseSQL("INSERT INTO insert_memory(time,device,value) VALUES(1,'d1','small')"); + Statement largeValue = + parseSQL( + "INSERT INTO insert_memory(time,device,value) VALUES(1,'d1','" + + "x".repeat(4096) + + "')"); + assertTrue(size > singleRow.ramBytesUsed()); + assertTrue(largeValue.ramBytesUsed() > singleRow.ramBytesUsed() + 4000); + + PreparedStatementInfo info = + PreparedStatementHelper.register(session, "insert_values", insert); + assertSame(insert, info.getSql()); + assertEquals(size, info.getMemorySizeInBytes()); + assertEquals(initialUsage + size, memory.getUsedMemoryInBytes()); + PreparedStatementHelper.unregister(session, "insert_values"); + assertNull(session.getPreparedStatement("insert_values")); + assertEquals(initialUsage, memory.getUsedMemoryInBytes()); + } finally { + if (session.getPreparedStatement("insert_values") != null) { + PreparedStatementHelper.unregister(session, "insert_values"); + } + cache.invalid("testdb"); + config.setDataNodeId(originalDataNodeId); + } + } + + @Test + public void testPreparedInsertValuesParseFailureMemoryAccounting() { + TsTable table = new TsTable("insert_failure_memory"); + table.addColumnSchema(new TimeColumnSchema("time", TSDataType.TIMESTAMP)); + table.addColumnSchema(new TagColumnSchema("device", TSDataType.STRING)); + table.addColumnSchema(new FieldColumnSchema("valid", TSDataType.INT32)); + table.addColumnSchema(new FieldColumnSchema("invalid", TSDataType.INT32)); + ITableCache cache = DataNodeTableCache.getInstance(); + cache.preUpdateTable("testdb", table, null); + cache.commitUpdateTable("testdb", "insert_failure_memory", null, null); + ClientSession session = new ClientSession(null); + IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + boolean originalPartialInsert = config.isEnablePartialInsert(); + int originalDataNodeId = config.getDataNodeId(); + config.setEnablePartialInsert(true); + config.setDataNodeId(0); + try { + InsertRows small = + (InsertRows) parseSQL("INSERT INTO insert_failure_memory VALUES(1,'d1',11,'bad')"); + InsertRows large = + (InsertRows) + parseSQL( + "INSERT INTO insert_failure_memory VALUES(1,'d1',11,'" + "x".repeat(4096) + "')"); + InsertRowStatement row = large.getInnerTreeStatement().getInsertRowStatementList().get(0); + assertEquals(1, row.getFailedMeasurementNumber()); + assertNull(row.getValues()[2]); + assertTrue(large.ramBytesUsed() > small.ramBytesUsed() + 4000); + + IMemoryBlock memory = Coordinator.getCoordinatorMemoryBlock(); + long initialUsage = memory.getUsedMemoryInBytes(); + PreparedStatementInfo info = + PreparedStatementHelper.register(session, "insert_failure", large); + assertEquals(large.ramBytesUsed(), info.getMemorySizeInBytes()); + assertEquals(initialUsage + large.ramBytesUsed(), memory.getUsedMemoryInBytes()); + PreparedStatementHelper.unregister(session, "insert_failure"); + assertEquals(initialUsage, memory.getUsedMemoryInBytes()); + } finally { + if (session.getPreparedStatement("insert_failure") != null) { + PreparedStatementHelper.unregister(session, "insert_failure"); + } + cache.invalid("testdb"); + config.setEnablePartialInsert(originalPartialInsert); + config.setDataNodeId(originalDataNodeId); + } + } + + @Test + public void testPreparedInsertValuesPreservesParseFailures() { + TsTable table = new TsTable("insert_partial"); + table.addColumnSchema(new TimeColumnSchema("time", TSDataType.TIMESTAMP)); + table.addColumnSchema(new TagColumnSchema("device", TSDataType.STRING)); + table.addColumnSchema(new FieldColumnSchema("valid", TSDataType.INT32)); + table.addColumnSchema(new FieldColumnSchema("invalid", TSDataType.INT32)); + ITableCache cache = DataNodeTableCache.getInstance(); + cache.preUpdateTable("testdb", table, null); + cache.commitUpdateTable("testdb", "insert_partial", null, null); + IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + boolean originalPartialInsert = config.isEnablePartialInsert(); + config.setEnablePartialInsert(true); + try { + Prepare prepare = + (Prepare) + parseSQL( + "PREPARE insert_partial FROM INSERT INTO insert_partial " + + "VALUES(1,'d1',11,'invalid integer')"); + InsertRows original = (InsertRows) prepare.getSql(); + InsertRowStatement originalRow = + original.getInnerTreeStatement().getInsertRowStatementList().get(0); + assertEquals(1, originalRow.getFailedMeasurementNumber()); + for (int i = 0; i < 2; i++) { + InsertRows execution = original.copyForPreparedExecution(); + InsertRowStatement executionRow = + execution.getInnerTreeStatement().getInsertRowStatementList().get(0); + execution.getInnerTreeStatement().semanticCheck(); + assertEquals(originalRow.getFailedMeasurements(), executionRow.getFailedMeasurements()); + assertEquals(originalRow.getFailedMessages(), executionRow.getFailedMessages()); + assertEquals(11, executionRow.getValues()[1]); + assertNotSame( + originalRow.getFailedMeasurementInfoMap(), executionRow.getFailedMeasurementInfoMap()); + executionRow.removeAllFailedMeasurementMarks(); + assertEquals(1, originalRow.getFailedMeasurementNumber()); + assertNull(originalRow.getMeasurements()[2]); + assertNull(originalRow.getValues()[2]); + } + } finally { + cache.invalid("testdb"); + config.setEnablePartialInsert(originalPartialInsert); + } + } + @Test public void testExecute() { Statement statement = parseSQL("EXECUTE myquery");