[BugFix] Fix mv rewrite bugs with IllegalStateException (backport #63655)

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

@@ -1236,6 +1236,10 @@ public Map<String, Range<PartitionKey>> getRangePartitionMap() {
             if (partitionName.startsWith(ExpressionRangePartitionInfo.SHADOW_PARTITION_PREFIX)) {
                 continue;
             }
+            // ensure partitionName is in nameToPartition
+            if (!nameToPartition.containsKey(partitionName)) {
+                continue;
+            }
             rangePartitionMap.put(partitionName, rangePartitionInfo.getRange(partitionId));
         }
         return rangePartitionMap;
@@ -1281,6 +1285,9 @@ public Map<String, PListCell> getListPartitionItems(Optional<List<Column>> selec
             if (partitionName.startsWith(ExpressionRangePartitionInfo.SHADOW_PARTITION_PREFIX)) {
                 continue;
             }
+            if (!nameToPartition.containsKey(partitionName)) {
+                continue;
+            }
             // one item
             List<LiteralExpr> literalValues = listPartitionInfo.getLiteralExprValues().get(partitionId);
             if (CollectionUtils.isNotEmpty(literalValues)) {

fe/fe-core/src/main/java/com/starrocks/scheduler/mv/MVPCTRefreshPartitioner.java

@@ -251,7 +251,7 @@ public static boolean isNonPartitionedMVNeedToRefresh(Map<Long, TableSnapshotInf
         return false;
     }
 
-    protected void dropPartition(Database db, MaterializedView materializedView, String mvPartitionName) {
+    public void dropPartition(Database db, MaterializedView materializedView, String mvPartitionName) {
         String dropPartitionName = materializedView.getPartition(mvPartitionName).getName();
         Locker locker = new Locker();
         if (!locker.tryLockTableWithIntensiveDbLock(db.getId(), materializedView.getId(), LockType.WRITE,

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/materialization/AggregatedMaterializedViewPushDownRewriter.java

@@ -477,6 +477,8 @@ public AggRewriteInfo visitLogicalTableScan(OptExpression optExpression, AggRewr
                     .setPartitionByColumns(groupBys)
                     .build();
             OptExpression optAggOp = OptExpression.create(newAggOp, optExpression);
+            // derive logical property for newly created agg node
+            MvUtils.deriveLogicalProperty(optAggOp);
 
             // rewrite by mv.
             OptExpression rewritten = doRewritePushDownAgg(mvRewriteContext, ctx, optAggOp, rule);

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/materialization/MaterializedViewRewriter.java

@@ -830,7 +830,7 @@ private boolean processJoinDerive(OptExpression queryExpr, OptExpression mvExpr)
         Map<Integer, ColumnRefSet> tableToJoinColumns = Maps.newHashMap();
         // first is from left, second is from right
         List<Pair<ColumnRefOperator, ColumnRefOperator>> joinColumnPairs = Lists.newArrayList();
-        ColumnRefSet leftColumns = queryExpr.inputAt(0).getOutputColumns();
+        ColumnRefSet leftColumns = MvUtils.getOutputColumns(queryExpr.inputAt(0));
         boolean isSupported = isSupportedPredicate(queryOnPredicate, materializationContext.getQueryRefFactory(),
                 leftColumns, tableToJoinColumns, joinColumnPairs);
         if (!isSupported) {
@@ -1218,7 +1218,7 @@ private ColumnRefOperator getColumnRef(String columnName, LogicalScanOperator sc
     private ScalarOperator collectMvPrunePredicate(MaterializationContext mvContext) {
         final OptExpression mvExpression = mvContext.getMvExpression();
         final Set<ScalarOperator> conjuncts = MvUtils.getAllValidPredicates(mvExpression);
-        final ColumnRefSet mvOutputColumnRefSet = mvExpression.getOutputColumns();
+        final ColumnRefSet mvOutputColumnRefSet = MvUtils.getOutputColumns(mvExpression);
         // conjuncts related to partition and distribution
         final List<ScalarOperator> mvPrunePredicates = Lists.newArrayList();
 
@@ -1738,7 +1738,7 @@ private PredicateSplit getUnionRewriteQueryCompensation(RewriteContext rewriteCo
                 .orElse(new ColumnRefSet());
 
         ColumnRefSet queryOutputColumnRefs = unionRewriteMode.isPullPredicateRewriteV2() ?
-                rewriteContext.getQueryExpression().getOutputColumns() : null;
+                MvUtils.getOutputColumns(rewriteContext.getQueryExpression()) : null;
         Set<ScalarOperator> queryExtraPredicates = queryPredicates.stream()
                 .filter(pred -> isPullUpQueryPredicate(pred, mvPredicateUsedColRefs,
                         queryOnPredicateUsedColRefs, queryOutputColumnRefs))
@@ -1750,7 +1750,7 @@ private PredicateSplit getUnionRewriteQueryCompensation(RewriteContext rewriteCo
         final ScalarOperator queryExtraPredicate = Utils.compoundAnd(queryExtraPredicates);
 
         // query's output should contain all the extra predicates otherwise it cannot be pulled then.
-        ColumnRefSet queryOutputColumnSet = rewriteContext.getQueryExpression().getOutputColumns();
+        ColumnRefSet queryOutputColumnSet = MvUtils.getOutputColumns(rewriteContext.getQueryExpression());
         if (!queryOutputColumnSet.containsAll(queryExtraPredicate.getUsedColumns())) {
             return null;
         }
@@ -2059,7 +2059,8 @@ private Operator doPruneEnforcedColumns(OptExpression queryExpr, Set<ColumnRefOp
                     .filter(entry -> !enforcedNonExistedColumns.contains(entry.getKey()))
                     .forEach(entry -> newColumnRefMap.put(entry.getKey(), entry.getValue()));
         } else {
-            queryExpr.getOutputColumns().getColumnRefOperators(materializationContext.getQueryRefFactory())
+            MvUtils.getOutputColumns(queryExpr)
+                    .getColumnRefOperators(materializationContext.getQueryRefFactory())
                     .stream()
                     .filter(column -> !enforcedNonExistedColumns.contains(column))
                     .forEach(column -> newColumnRefMap.put(column, column));

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/materialization/MvUtils.java

@@ -1678,4 +1678,15 @@ public static String shrinkToSize(String input, int maxLength) {
         }
         return input.length() > maxLength ? input.substring(0, maxLength) : input;
     }
+
+    /**
+     * Ensure the logical property of the given opt expression is derived, and return its output columns.
+     * Ensure optExpression has derived logical property, otherwise OptExpression.getOutputColumns will fail.
+     */
+    public static ColumnRefSet getOutputColumns(OptExpression optExpression) {
+        if (optExpression.getLogicalProperty() == null) {
+            deriveLogicalProperty(optExpression);
+        }
+        return optExpression.getOutputColumns();
+    }
 }