Fix database time partition config cache

Author: zerolbsony

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/CNPhysicalPlanGeneratorTest.java

@@ -335,7 +335,7 @@ public void databaseWithoutTemplateGeneratorTest() throws Exception {
       tDatabaseSchema.setTTL(i + 1);
       tDatabaseSchema.setDataReplicationFactor(i);
       tDatabaseSchema.setSchemaReplicationFactor(i);
-      tDatabaseSchema.setTimePartitionInterval(i);
+      tDatabaseSchema.setTimePartitionInterval(i + 1);
       DatabaseSchemaPlan databaseSchemaPlan =
           new DatabaseSchemaPlan(ConfigPhysicalPlanType.CreateDatabase, tDatabaseSchema);
       clusterSchemaInfo.createDatabase(databaseSchemaPlan);
@@ -453,7 +453,7 @@ public void templateAndDatabaseCompletedTest() throws Exception {
       tDatabaseSchema.setTTL(i + 1);
       tDatabaseSchema.setDataReplicationFactor(i);
       tDatabaseSchema.setSchemaReplicationFactor(i);
-      tDatabaseSchema.setTimePartitionInterval(i);
+      tDatabaseSchema.setTimePartitionInterval(i + 1);
       final DatabaseSchemaPlan databaseSchemaPlan =
           new DatabaseSchemaPlan(ConfigPhysicalPlanType.CreateDatabase, tDatabaseSchema);
       clusterSchemaInfo.createDatabase(databaseSchemaPlan);

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

@@ -133,43 +133,43 @@ public static List<DataPartitionQueryParam> computeTableDataPartitionParams(
     String database = getDatabaseName(statement, context);
     if (statement instanceof InsertTabletStatement) {
       final InsertTabletStatement insertTabletStatement = (InsertTabletStatement) statement;
-      final Map<IDeviceID, Set<TTimePartitionSlot>> timePartitionSlotMap = new HashMap<>();
+      final Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
       for (int i = 0; i < insertTabletStatement.getRowCount(); i++) {
-        timePartitionSlotMap
-            .computeIfAbsent(insertTabletStatement.getTableDeviceID(i), id -> new HashSet<>())
-            .add(insertTabletStatement.getTimePartitionSlot(i, database));
+        dataPartitionQueryParamMap
+            .computeIfAbsent(insertTabletStatement.getTableDeviceID(i), id -> new ArrayList<>())
+            .add(insertTabletStatement.getTimes()[i]);
       }
-      return computeDataPartitionParams(timePartitionSlotMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     } else if (statement instanceof InsertMultiTabletsStatement) {
       final InsertMultiTabletsStatement insertMultiTabletsStatement =
           (InsertMultiTabletsStatement) statement;
-      final Map<IDeviceID, Set<TTimePartitionSlot>> timePartitionSlotMap = new HashMap<>();
+      final Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
       for (final InsertTabletStatement insertTabletStatement :
           insertMultiTabletsStatement.getInsertTabletStatementList()) {
         for (int i = 0; i < insertTabletStatement.getRowCount(); i++) {
-          timePartitionSlotMap
-              .computeIfAbsent(insertTabletStatement.getTableDeviceID(i), id -> new HashSet<>())
-              .add(insertTabletStatement.getTimePartitionSlot(i, database));
+          dataPartitionQueryParamMap
+              .computeIfAbsent(insertTabletStatement.getTableDeviceID(i), id -> new ArrayList<>())
+              .add(insertTabletStatement.getTimes()[i]);
         }
       }
-      return computeDataPartitionParams(timePartitionSlotMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     } else if (statement instanceof InsertRowStatement) {
       final InsertRowStatement insertRowStatement = (InsertRowStatement) statement;
-      return computeDataPartitionParams(
+      return computeDataPartitionParamsByTime(
           Collections.singletonMap(
               insertRowStatement.getTableDeviceID(),
-              Collections.singleton(insertRowStatement.getTimePartitionSlot(database))),
+              Collections.singletonList(insertRowStatement.getTime())),
           database);
     } else if (statement instanceof InsertRowsStatement) {
       final InsertRowsStatement insertRowsStatement = (InsertRowsStatement) statement;
-      final Map<IDeviceID, Set<TTimePartitionSlot>> timePartitionSlotMap = new HashMap<>();
+      final Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
       for (final InsertRowStatement insertRowStatement :
           insertRowsStatement.getInsertRowStatementList()) {
-        timePartitionSlotMap
-            .computeIfAbsent(insertRowStatement.getTableDeviceID(), id -> new HashSet<>())
-            .add(insertRowStatement.getTimePartitionSlot(database));
+        dataPartitionQueryParamMap
+            .computeIfAbsent(insertRowStatement.getTableDeviceID(), id -> new ArrayList<>())
+            .add(insertRowStatement.getTime());
       }
-      return computeDataPartitionParams(timePartitionSlotMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     }
     throw new UnsupportedOperationException(
         DataNodeQueryMessages.COMPUTEDATAPARTITIONPARAMS_FOR + statement);
@@ -184,95 +184,51 @@ public static List<DataPartitionQueryParam> computeTreeDataPartitionParams(
       return Collections.singletonList(dataPartitionQueryParam);
     } else if (statement instanceof InsertRowStatement) {
       InsertRowStatement insertRowStatement = (InsertRowStatement) statement;
-      if (database == null) {
-        DataPartitionQueryParam dataPartitionQueryParam = new DataPartitionQueryParam();
-        dataPartitionQueryParam.setDeviceID(
-            insertRowStatement.getDevicePath().getIDeviceIDAsFullDevice());
-        dataPartitionQueryParam.setTimeList(
-            Collections.singletonList(insertRowStatement.getTime()));
-        return Collections.singletonList(dataPartitionQueryParam);
-      }
-      return computeDataPartitionParams(
+      return computeDataPartitionParamsByTime(
           Collections.singletonMap(
               insertRowStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-              Collections.singleton(insertRowStatement.getTimePartitionSlot(database))),
+              Collections.singletonList(insertRowStatement.getTime())),
           database);
     } else if (statement instanceof InsertMultiTabletsStatement) {
       InsertMultiTabletsStatement insertMultiTabletsStatement =
           (InsertMultiTabletsStatement) statement;
-      if (database == null) {
-        Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
-        for (InsertTabletStatement insertTabletStatement :
-            insertMultiTabletsStatement.getInsertTabletStatementList()) {
-          List<Long> timeList =
-              dataPartitionQueryParamMap.computeIfAbsent(
-                  insertTabletStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-                  k -> new ArrayList<>());
-          for (int i = 0; i < insertTabletStatement.getRowCount(); i++) {
-            timeList.add(insertTabletStatement.getTimes()[i]);
-          }
-        }
-        return computeDataPartitionParamsByTime(dataPartitionQueryParamMap);
-      }
-      Map<IDeviceID, Set<TTimePartitionSlot>> dataPartitionQueryParamMap = new HashMap<>();
+      Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
       for (InsertTabletStatement insertTabletStatement :
           insertMultiTabletsStatement.getInsertTabletStatementList()) {
-        Set<TTimePartitionSlot> timePartitionSlotSet =
+        List<Long> timeList =
             dataPartitionQueryParamMap.computeIfAbsent(
                 insertTabletStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-                k -> new HashSet<>());
-        timePartitionSlotSet.addAll(insertTabletStatement.getTimePartitionSlots(database));
+                k -> new ArrayList<>());
+        for (int i = 0; i < insertTabletStatement.getRowCount(); i++) {
+          timeList.add(insertTabletStatement.getTimes()[i]);
+        }
       }
-      return computeDataPartitionParams(dataPartitionQueryParamMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     } else if (statement instanceof InsertRowsOfOneDeviceStatement) {
       final InsertRowsOfOneDeviceStatement insertRowsOfOneDeviceStatement =
           (InsertRowsOfOneDeviceStatement) statement;
-      if (database == null) {
-        Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
-        List<Long> timeList =
-            dataPartitionQueryParamMap.computeIfAbsent(
-                insertRowsOfOneDeviceStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-                k -> new ArrayList<>());
-        for (InsertRowStatement insertRowStatement :
-            insertRowsOfOneDeviceStatement.getInsertRowStatementList()) {
-          timeList.add(insertRowStatement.getTime());
-        }
-        return computeDataPartitionParamsByTime(dataPartitionQueryParamMap);
-      }
-      Map<IDeviceID, Set<TTimePartitionSlot>> dataPartitionQueryParamMap = new HashMap<>();
-      Set<TTimePartitionSlot> timePartitionSlotSet =
+      Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
+      List<Long> timeList =
           dataPartitionQueryParamMap.computeIfAbsent(
               insertRowsOfOneDeviceStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-              k -> new HashSet<>());
+              k -> new ArrayList<>());
       for (InsertRowStatement insertRowStatement :
           insertRowsOfOneDeviceStatement.getInsertRowStatementList()) {
-        timePartitionSlotSet.add(insertRowStatement.getTimePartitionSlot(database));
+        timeList.add(insertRowStatement.getTime());
       }
-      return computeDataPartitionParams(dataPartitionQueryParamMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     } else if (statement instanceof InsertRowsStatement) {
       final InsertRowsStatement insertRowsStatement = (InsertRowsStatement) statement;
-      if (database == null) {
-        Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
-        for (InsertRowStatement insertRowStatement :
-            insertRowsStatement.getInsertRowStatementList()) {
-          dataPartitionQueryParamMap
-              .computeIfAbsent(
-                  insertRowStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-                  k -> new ArrayList<>())
-              .add(insertRowStatement.getTime());
-        }
-        return computeDataPartitionParamsByTime(dataPartitionQueryParamMap);
-      }
-      Map<IDeviceID, Set<TTimePartitionSlot>> dataPartitionQueryParamMap = new HashMap<>();
+      Map<IDeviceID, List<Long>> dataPartitionQueryParamMap = new HashMap<>();
       for (InsertRowStatement insertRowStatement :
           insertRowsStatement.getInsertRowStatementList()) {
-        Set<TTimePartitionSlot> timePartitionSlotSet =
-            dataPartitionQueryParamMap.computeIfAbsent(
+        dataPartitionQueryParamMap
+            .computeIfAbsent(
                 insertRowStatement.getDevicePath().getIDeviceIDAsFullDevice(),
-                k -> new HashSet<>());
-        timePartitionSlotSet.add(insertRowStatement.getTimePartitionSlot(database));
+                k -> new ArrayList<>())
+            .add(insertRowStatement.getTime());
       }
-      return computeDataPartitionParams(dataPartitionQueryParamMap, database);
+      return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, database);
     }
     throw new UnsupportedOperationException(
         DataNodeQueryMessages.COMPUTEDATAPARTITIONPARAMS_FOR + statement);
@@ -282,14 +238,12 @@ private static DataPartitionQueryParam getTreeDataPartitionQueryParam(
       InsertTabletStatement statement, String database) {
     DataPartitionQueryParam dataPartitionQueryParam = new DataPartitionQueryParam();
     dataPartitionQueryParam.setDeviceID(statement.getDevicePath().getIDeviceIDAsFullDevice());
-    if (database == null) {
-      List<Long> timeList = new ArrayList<>();
-      for (int i = 0; i < statement.getRowCount(); i++) {
-        timeList.add(statement.getTimes()[i]);
-      }
-      dataPartitionQueryParam.setTimeList(timeList);
-    } else {
-      dataPartitionQueryParam.setTimePartitionSlotList(statement.getTimePartitionSlots(database));
+    List<Long> timeList = new ArrayList<>();
+    for (int i = 0; i < statement.getRowCount(); i++) {
+      timeList.add(statement.getTimes()[i]);
+    }
+    dataPartitionQueryParam.setTimeList(timeList);
+    if (database != null) {
       dataPartitionQueryParam.setDatabaseName(database);
     }
     return dataPartitionQueryParam;
@@ -316,11 +270,19 @@ public static List<DataPartitionQueryParam> computeDataPartitionParams(
 
   private static List<DataPartitionQueryParam> computeDataPartitionParamsByTime(
       Map<IDeviceID, List<Long>> dataPartitionQueryParamMap) {
+    return computeDataPartitionParamsByTime(dataPartitionQueryParamMap, null);
+  }
+
+  private static List<DataPartitionQueryParam> computeDataPartitionParamsByTime(
+      Map<IDeviceID, List<Long>> dataPartitionQueryParamMap, String databaseName) {
     List<DataPartitionQueryParam> dataPartitionQueryParams = new ArrayList<>();
     for (Map.Entry<IDeviceID, List<Long>> entry : dataPartitionQueryParamMap.entrySet()) {
       DataPartitionQueryParam dataPartitionQueryParam = new DataPartitionQueryParam();
       dataPartitionQueryParam.setDeviceID(entry.getKey());
       dataPartitionQueryParam.setTimeList(entry.getValue());
+      if (databaseName != null) {
+        dataPartitionQueryParam.setDatabaseName(databaseName);
+      }
       dataPartitionQueryParams.add(dataPartitionQueryParam);
     }
     return dataPartitionQueryParams;

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

@@ -34,6 +34,7 @@
 import org.apache.iotdb.commons.partition.executor.SeriesPartitionExecutor;
 import org.apache.iotdb.commons.path.PathPatternTree;
 import org.apache.iotdb.commons.schema.SchemaConstant;
+import org.apache.iotdb.commons.utils.PathUtils;
 import org.apache.iotdb.commons.utils.TimePartitionUtils;
 import org.apache.iotdb.confignode.rpc.thrift.TDataPartitionReq;
 import org.apache.iotdb.confignode.rpc.thrift.TDataPartitionTableResp;
@@ -335,23 +336,36 @@ public void ensureDatabaseTimePartitionConfig(Set<String> databases) {
     if (missingDatabases.isEmpty()) {
       return;
     }
+    final boolean needFetchTableModelDatabase =
+        missingDatabases.stream().anyMatch(PathUtils::isTableModelDatabase);
+    final boolean needFetchTreeModelDatabase =
+        missingDatabases.stream().anyMatch(database -> !PathUtils.isTableModelDatabase(database));
     try (final ConfigNodeClient client =
         configNodeClientManager.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
-      final TDatabaseSchemaResp databaseSchemaResp =
-          client.getMatchedDatabaseSchemas(
-              new TGetDatabaseReq(
-                      Arrays.asList("root", "**"), SchemaConstant.ALL_MATCH_SCOPE_BINARY)
-                  .setIsTableModel(false)
-                  .setCanSeeAuditDB(true));
-      if (databaseSchemaResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-        partitionCache.updateDatabaseCache(databaseSchemaResp.getDatabaseSchemaMap());
+      if (needFetchTreeModelDatabase) {
+        updateDatabaseCache(client, false);
+      }
+      if (needFetchTableModelDatabase) {
+        updateDatabaseCache(client, true);
       }
     } catch (final ClientManagerException | TException e) {
       throw new StatementAnalyzeException(
           "An error occurred when fetching database time partition config:" + e.getMessage());
     }
   }
 
+  private void updateDatabaseCache(final ConfigNodeClient client, final boolean isTableModel)
+      throws TException {
+    final TDatabaseSchemaResp databaseSchemaResp =
+        client.getMatchedDatabaseSchemas(
+            new TGetDatabaseReq(Arrays.asList("root", "**"), SchemaConstant.ALL_MATCH_SCOPE_BINARY)
+                .setIsTableModel(isTableModel)
+                .setCanSeeAuditDB(true));
+    if (databaseSchemaResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+      partitionCache.updateDatabaseCache(databaseSchemaResp.getDatabaseSchemaMap());
+    }
+  }
+
   @Override
   public SchemaPartition getOrCreateSchemaPartition(
       final String database, final List<IDeviceID> deviceIDs, final String userName) {
@@ -436,6 +450,10 @@ private Map<String, List<DataPartitionQueryParam>> splitDataPartitionQueryParam(
         database = dataPartitionQueryParam.getDatabaseName();
       }
       if (database != null) {
+        if (PathUtils.isTableModelDatabase(database)) {
+          partitionCache.checkAndAutoCreateDatabase(database, isAutoCreate, userName);
+        }
+        ensureDatabaseTimePartitionConfig(Collections.singleton(database));
         final String finalDatabase = database;
         if (dataPartitionQueryParam.hasTimeList()) {
           dataPartitionQueryParam.setTimePartitionSlotList(

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/partition/PartitionCache.java

@@ -366,8 +366,9 @@ private void createDatabaseAndUpdateCache(
             throw new IoTDBRuntimeException(tsStatus.message, tsStatus.code);
           }
         }
-        // Try to update database cache when all databases have already been created
-        if (needRefreshCacheFromConfigNode) {
+        // Fetch the completed schema from ConfigNode after auto-creation. The locally constructed
+        // schema may miss default time partition settings that ConfigNode fills in.
+        if (needRefreshCacheFromConfigNode || !successFullyCreatedDatabase.isEmpty()) {
           fetchDatabaseAndUpdateCache(client, false);
         } else {
           updateDatabaseCache(successFullyCreatedDatabaseSchema);
@@ -404,8 +405,9 @@ private void createDatabaseAndUpdateCache(final String database, final String us
       databaseSchema.setIsTableModel(true);
       final TSStatus tsStatus = client.setDatabase(databaseSchema);
       if (TSStatusCode.SUCCESS_STATUS.getStatusCode() == tsStatus.getCode()) {
-        // Try to update cache by databases successfully created
-        updateDatabaseCache(Collections.singletonMap(database, databaseSchema));
+        // Fetch the completed schema from ConfigNode after auto-creation. The locally constructed
+        // schema may miss default time partition settings that ConfigNode fills in.
+        fetchDatabaseAndUpdateCache(client, true);
       } else if (TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode() == tsStatus.getCode()) {
         fetchDatabaseAndUpdateCache(client, true);
       } else {