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

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

@@ -653,6 +653,49 @@ public Object clone() {
  return info;
  }
 
+ @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();

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

@@ -829,141 +829,46 @@ 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());
  }
 
- // reset partition info and idToPartition map
- if (partitionInfo.isRangePartition()) {
- RangePartitionInfo rangePartitionInfo = (RangePartitionInfo) partitionInfo;
- RangePartitionInfo origRangePartitionInfo = (RangePartitionInfo) rangePartitionInfo.clone();
- for (Long partitionId : origPartNameToId.values()) {
- rangePartitionInfo.dropPartition(partitionId);
- }
- Map<Long, Partition> origIdToPartition = Maps.newHashMap(idToPartition);
- idToPartition.clear();
- physicalPartitionIdToPartitionId.clear();
- physicalPartitionNameToPartitionId.clear();
- for (Map.Entry<String, Long> entry : origPartNameToId.entrySet()) {
- long newPartId = globalStateMgr.getNextId();
- // preserve existing info
- DataProperty dataProperty = origRangePartitionInfo.getDataProperty(entry.getValue());
- boolean inMemory = origRangePartitionInfo.getIsInMemory(entry.getValue());
- DataCacheInfo dataCacheInfo = origRangePartitionInfo.getDataCacheInfo(entry.getValue());
- Range<PartitionKey> range = origRangePartitionInfo.getIdToRange(false).get(entry.getValue());
- // replace with new info
- rangePartitionInfo.addPartition(newPartId, false, range, dataProperty, (short) restoreReplicationNum,
- inMemory, dataCacheInfo);
- idToPartition.put(newPartId, origIdToPartition.get(entry.getValue()));
- Partition partition = idToPartition.get(newPartId);
- partition.setIdForRestore(newPartId);
- List<PhysicalPartition> origPhysicalPartitions = Lists.newArrayList(partition.getSubPartitions());
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- partition.removeSubPartition(physicalPartition.getId());
- }
- });
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- physicalPartition.setIdForRestore(globalStateMgr.getNextId());
- physicalPartition.setParentId(newPartId);
- partition.addSubPartition(physicalPartition);
- }
- physicalPartitionIdToPartitionId.put(physicalPartition.getId(), newPartId);
- physicalPartitionNameToPartitionId.put(physicalPartition.getName(), newPartId);
- });
- }
- } else if (partitionInfo.isListPartition()) {
- ListPartitionInfo listPartitionInfo = (ListPartitionInfo) partitionInfo;
- ListPartitionInfo origListPartitionInfo = (ListPartitionInfo) listPartitionInfo.clone();
- for (Long partitionId : origPartNameToId.values()) {
- listPartitionInfo.dropPartition(partitionId);
- }
- Map<Long, Partition> origIdToPartition = Maps.newHashMap(idToPartition);
- idToPartition.clear();
- physicalPartitionIdToPartitionId.clear();
- physicalPartitionNameToPartitionId.clear();
- for (Map.Entry<String, Long> entry : origPartNameToId.entrySet()) {
- long newPartId = globalStateMgr.getNextId();
- // preserve existing info
- DataProperty dataProperty = origListPartitionInfo.getDataProperty(entry.getValue());
- boolean inMemory = origListPartitionInfo.getIsInMemory(entry.getValue());
- DataCacheInfo dataCacheInfo = origListPartitionInfo.getDataCacheInfo(entry.getValue());
- List<String> values = origListPartitionInfo.getIdToValues().get(entry.getValue());
- List<List<String>> multiValues = origListPartitionInfo.getIdToMultiValues().get(entry.getValue());
- // replace with new info
- try {
- listPartitionInfo.addPartition(idToColumn, newPartId, dataProperty, (short) restoreReplicationNum,
- inMemory, dataCacheInfo, values, multiValues);
- } catch (AnalysisException e) {
- return new Status(ErrCode.COMMON_ERROR, "Failed to add partition " + e.getMessage());
+ // 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());
  }
- idToPartition.put(newPartId, origIdToPartition.get(entry.getValue()));
- Partition partition = idToPartition.get(newPartId);
- partition.setIdForRestore(newPartId);
- List<PhysicalPartition> origPhysicalPartitions = Lists.newArrayList(partition.getSubPartitions());
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- partition.removeSubPartition(physicalPartition.getId());
- }
- });
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- physicalPartition.setIdForRestore(globalStateMgr.getNextId());
- physicalPartition.setParentId(newPartId);
- partition.addSubPartition(physicalPartition);
- }
- physicalPartitionIdToPartitionId.put(physicalPartition.getId(), newPartId);
- physicalPartitionNameToPartitionId.put(physicalPartition.getName(), newPartId);
- });
- }
- } else if (partitionInfo.isUnPartitioned()) {
- // Single partitioned
- PartitionInfo origPartitionInfo = (PartitionInfo) partitionInfo.clone();
- for (Long partitionId : origPartNameToId.values()) {
- partitionInfo.dropPartition(partitionId);
- }
- Map<Long, Partition> origIdToPartition = Maps.newHashMap(idToPartition);
- idToPartition.clear();
- physicalPartitionIdToPartitionId.clear();
- physicalPartitionNameToPartitionId.clear();
- for (Map.Entry<String, Long> entry : origPartNameToId.entrySet()) {
- long newPartId = globalStateMgr.getNextId();
- DataProperty dataProperty = origPartitionInfo.getDataProperty(entry.getValue());
- boolean inMemory = origPartitionInfo.getIsInMemory(entry.getValue());
- DataCacheInfo dataCacheInfo = origPartitionInfo.getDataCacheInfo(entry.getValue());
- partitionInfo.addPartition(newPartId, dataProperty, (short) restoreReplicationNum, inMemory,
- dataCacheInfo);
- idToPartition.put(newPartId, origIdToPartition.get(entry.getValue()));
- Partition partition = idToPartition.get(newPartId);
- partition.setIdForRestore(newPartId);
- List<PhysicalPartition> origPhysicalPartitions = Lists.newArrayList(partition.getSubPartitions());
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- partition.removeSubPartition(physicalPartition.getId());
- }
- });
- origPhysicalPartitions.forEach(physicalPartition -> {
- if (physicalPartition.getId() != newPartId) {
- physicalPartition.setIdForRestore(globalStateMgr.getNextId());
- physicalPartition.setParentId(newPartId);
- partition.addSubPartition(physicalPartition);
- }
- physicalPartitionIdToPartitionId.put(physicalPartition.getId(), newPartId);
- physicalPartitionNameToPartitionId.put(physicalPartition.getName(), newPartId);
- });
- }
- } else {
- return new Status(ErrCode.UNSUPPORTED, "Unsupported partition type: " + partitionInfo.getType());
+ });
+ 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);
+ });
  }
 
  // reset replication number for olaptable

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

@@ -311,4 +311,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

@@ -569,4 +569,29 @@ protected Object clone() {
  info.isMultiColumnPartition = partitionColumnIds.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);
+ }
+ }
+ }
 }