[Refactor] Refactor reset ids for restore (#52075)

Signed-off-by: xiangguangyxg <[email protected]>
(cherry picked from commit 9db7976)

# Conflicts:
#	fe/fe-core/src/main/java/com/starrocks/catalog/ListPartitionInfo.java
#	fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java

fe/fe-core/src/main/java/com/starrocks/catalog/ListPartitionInfo.java

@@ -513,4 +513,58 @@ public Object clone() {
  info.automaticPartition = this.automaticPartition;
  return info;
  }
+<<<<<<< HEAD
+=======
+
+ @Override
+ public void setPartitionIdsForRestore(Map<Long, Long> partitionOldIdToNewId) {
+ super.setPartitionIdsForRestore(partitionOldIdToNewId);
+
+ Map<Long, List<List<String>>> oldIdToMultiValues = this.idToMultiValues;
+ Map<Long, List<List<LiteralExpr>>> oldIdToMultiLiteralExprValues = this.idToMultiLiteralExprValues;
+ Map<Long, List<String>> oldIdToValues = this.idToValues;
+ Map<Long, List<LiteralExpr>> oldIdToLiteralExprValues = this.idToLiteralExprValues;
+ Map<Long, Boolean> oldIdToIsTempPartition = this.idToIsTempPartition;
+
+ this.idToMultiValues = new HashMap<>();
+ this.idToMultiLiteralExprValues = new HashMap<>();
+ this.idToValues = new HashMap<>();
+ this.idToLiteralExprValues = new HashMap<>();
+ this.idToIsTempPartition = new HashMap<>();
+
+ for (Map.Entry<Long, Long> entry : partitionOldIdToNewId.entrySet()) {
+ Long oldId = entry.getKey();
+ Long newId = entry.getValue();
+
+ List<List<String>> multiValues = oldIdToMultiValues.get(oldId);
+ if (multiValues != null) {
+ this.idToMultiValues.put(newId, multiValues);
+ }
+ List<List<LiteralExpr>> multiLiteralExprValues = oldIdToMultiLiteralExprValues.get(oldId);
+ if (multiLiteralExprValues != null) {
+ this.idToMultiLiteralExprValues.put(newId, multiLiteralExprValues);
+ }
+ List<String> values = oldIdToValues.get(oldId);
+ if (values != null) {
+ this.idToValues.put(newId, values);
+ }
+ List<LiteralExpr> literalExprValues = oldIdToLiteralExprValues.get(oldId);
+ if (literalExprValues != null) {
+ this.idToLiteralExprValues.put(newId, literalExprValues);
+ }
+ Boolean isTempPartition = oldIdToIsTempPartition.get(oldId);
+ if (isTempPartition != null) {
+ this.idToIsTempPartition.put(newId, isTempPartition);
+ }
+ }
+ }
+
+ @Override
+ public void gsonPostProcess() throws IOException {
+ super.gsonPostProcess();
+ if (partitionColumnIds.size() <= 0) {
+ partitionColumnIds = deprecatedColumns.stream().map(Column::getColumnId).collect(Collectors.toList());
+ }
+ }
+>>>>>>> 9db79768de ([Refactor] Refactor reset ids for restore (#52075))
 }

fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java

@@ -739,19 +739,20 @@ public Status resetIdsForRestore(GlobalStateMgr globalStateMgr, Database db, int
  // base index
  baseIndexId = newIdxId;
  }
- indexIdToMeta.put(newIdxId, origIndexIdToMeta.get(entry.getKey()));
- indexIdToMeta.get(newIdxId).setIndexIdForRestore(newIdxId);
- indexIdToMeta.get(newIdxId).setSchemaId(newIdxId);
+ MaterializedIndexMeta indexMeta = origIndexIdToMeta.get(entry.getKey());
+ indexMeta.setIndexIdForRestore(newIdxId);
+ indexMeta.setSchemaId(newIdxId);
+ indexIdToMeta.put(newIdxId, indexMeta);
  indexNameToId.put(entry.getValue(), newIdxId);
  }
 
- // generate a partition name to id map
- Map<String, Long> origPartNameToId = Maps.newHashMap();
- for (Partition partition : idToPartition.values()) {
- origPartNameToId.put(partition.getName(), partition.getId());
- LOG.info("partition id {} sub partition {}", partition.getId(), partition.getSubPartitions());
+ // generate a partition old id to new id map
+ Map<Long, Long> partitionOldIdToNewId = Maps.newHashMap();
+ for (Long id : idToPartition.keySet()) {
+ partitionOldIdToNewId.put(id, globalStateMgr.getNextId());
  }
 
+<<<<<<< HEAD
  // reset partition info and idToPartition map
  if (partitionInfo.isRangePartition()) {
  RangePartitionInfo rangePartitionInfo = (RangePartitionInfo) partitionInfo;
@@ -868,6 +869,35 @@ public Status resetIdsForRestore(GlobalStateMgr globalStateMgr, Database db, int
  }
  } else {
  return new Status(ErrCode.UNSUPPORTED, "Unsupported partition type: " + partitionInfo.getType());
+=======
+ // reset partiton info
+ partitionInfo.setPartitionIdsForRestore(partitionOldIdToNewId);
+
+ // reset partitions
+ List<Partition> partitions = Lists.newArrayList(idToPartition.values());
+ idToPartition.clear();
+ physicalPartitionIdToPartitionId.clear();
+ physicalPartitionNameToPartitionId.clear();
+ for (Partition partition : partitions) {
+ long newPartitionId = partitionOldIdToNewId.get(partition.getId());
+ partition.setIdForRestore(newPartitionId);
+ idToPartition.put(newPartitionId, partition);
+ List<PhysicalPartition> origPhysicalPartitions = Lists.newArrayList(partition.getSubPartitions());
+ origPhysicalPartitions.forEach(physicalPartition -> {
+ if (physicalPartition.getId() != newPartitionId) {
+ partition.removeSubPartition(physicalPartition.getId());
+ }
+ });
+ origPhysicalPartitions.forEach(physicalPartition -> {
+ if (physicalPartition.getId() != newPartitionId) {
+ physicalPartition.setIdForRestore(globalStateMgr.getNextId());
+ physicalPartition.setParentId(newPartitionId);
+ partition.addSubPartition(physicalPartition);
+ }
+ physicalPartitionIdToPartitionId.put(physicalPartition.getId(), newPartitionId);
+ physicalPartitionNameToPartitionId.put(physicalPartition.getName(), newPartitionId);
+ });
+>>>>>>> 9db79768de ([Refactor] Refactor reset ids for restore (#52075))
  }
 
  // reset replication number for olaptable

fe/fe-core/src/main/java/com/starrocks/catalog/PartitionInfo.java

@@ -337,4 +337,44 @@ protected Object clone() {
  throw new RuntimeException(e);
  }
  }
+
+ public void setPartitionIdsForRestore(Map<Long, Long> partitionOldIdToNewId) {
+ Map<Long, DataProperty> oldIdToDataProperty = this.idToDataProperty;
+ Map<Long, Short> oldIdToReplicationNum = this.idToReplicationNum;
+ Map<Long, Boolean> oldIdToInMemory = this.idToInMemory;
+ Map<Long, TTabletType> oldIdToTabletType = this.idToTabletType;
+ Map<Long, DataCacheInfo> oldIdToStorageCacheInfo = this.idToStorageCacheInfo;
+
+ this.idToDataProperty = new HashMap<>();
+ this.idToReplicationNum = new HashMap<>();
+ this.idToInMemory = new HashMap<>();
+ this.idToTabletType = new HashMap<>();
+ this.idToStorageCacheInfo = new HashMap<>();
+
+ for (Map.Entry<Long, Long> entry : partitionOldIdToNewId.entrySet()) {
+ Long oldId = entry.getKey();
+ Long newId = entry.getValue();
+
+ DataProperty dataProperty = oldIdToDataProperty.get(oldId);
+ if (dataProperty != null) {
+ this.idToDataProperty.put(newId, dataProperty);
+ }
+ Short replicationNum = oldIdToReplicationNum.get(oldId);
+ if (replicationNum != null) {
+ this.idToReplicationNum.put(newId, replicationNum);
+ }
+ Boolean inMemory = oldIdToInMemory.get(oldId);
+ if (inMemory != null) {
+ this.idToInMemory.put(newId, inMemory);
+ }
+ TTabletType tabletType = oldIdToTabletType.get(oldId);
+ if (tabletType != null) {
+ this.idToTabletType.put(newId, tabletType);
+ }
+ DataCacheInfo dataCacheInfo = oldIdToStorageCacheInfo.get(oldId);
+ if (dataCacheInfo != null) {
+ this.idToStorageCacheInfo.put(newId, dataCacheInfo);
+ }
+ }
+ }
 }

fe/fe-core/src/main/java/com/starrocks/catalog/RangePartitionInfo.java

@@ -589,5 +589,30 @@ protected Object clone() {
  info.isMultiColumnPartition = partitionColumns.size() > 1;
  return info;
  }
+
+ @Override
+ public void setPartitionIdsForRestore(Map<Long, Long> partitionOldIdToNewId) {
+ super.setPartitionIdsForRestore(partitionOldIdToNewId);
+
+ Map<Long, Range<PartitionKey>> oldIdToRange = this.idToRange;
+ Map<Long, Range<PartitionKey>> oldIdToTempRange = this.idToTempRange;
+
+ this.idToRange = new ConcurrentHashMap<>();
+ this.idToTempRange = new ConcurrentHashMap<>();
+
+ for (Map.Entry<Long, Long> entry : partitionOldIdToNewId.entrySet()) {
+ Long oldId = entry.getKey();
+ Long newId = entry.getValue();
+
+ Range<PartitionKey> range = oldIdToRange.get(oldId);
+ if (range != null) {
+ this.idToRange.put(newId, range);
+ }
+ Range<PartitionKey> tempRange = oldIdToTempRange.get(oldId);
+ if (tempRange != null) {
+ this.idToTempRange.put(newId, tempRange);
+ }
+ }
+ }
 }