address review comments.

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/mv/AbstractSelectMaterializedIndexRule.java

@@ -50,6 +50,17 @@
  * Base class for selecting materialized index rules.
  */
 public abstract class AbstractSelectMaterializedIndexRule {
+ protected boolean shouldSelectIndex(LogicalOlapScan scan) {
+ switch (scan.getTable().getKeysType()) {
+ case AGG_KEYS:
+ case UNIQUE_KEYS:
+ case DUP_KEYS:
+ return !scan.isIndexSelected();
+ default:
+ return false;
+ }
+ }
+
  /**
  * 1. indexes have all the required columns.
  * 2. find matching key prefix most.

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/mv/SelectMaterializedIndexWithAggregate.java

@@ -74,7 +74,7 @@ public List<Rule> buildRules() {
  return ImmutableList.of(
  // only agg above scan
  // Aggregate(Scan)
- logicalAggregate(logicalOlapScan().when(LogicalOlapScan::shouldSelectIndex)).then(agg -> {
+ logicalAggregate(logicalOlapScan().when(this::shouldSelectIndex)).then(agg -> {
  LogicalOlapScan scan = agg.child();
  Pair<PreAggStatus, List<Long>> result = select(
  scan,
@@ -89,7 +89,7 @@ public List<Rule> buildRules() {
 
  // filter could push down scan.
  // Aggregate(Filter(Scan))
- logicalAggregate(logicalFilter(logicalOlapScan().when(LogicalOlapScan::shouldSelectIndex)))
+ logicalAggregate(logicalFilter(logicalOlapScan().when(this::shouldSelectIndex)))
  .then(agg -> {
  LogicalFilter<LogicalOlapScan> filter = agg.child();
  LogicalOlapScan scan = filter.child();
@@ -112,7 +112,7 @@ public List<Rule> buildRules() {
 
  // column pruning or other projections such as alias, etc.
  // Aggregate(Project(Scan))
- logicalAggregate(logicalProject(logicalOlapScan().when(LogicalOlapScan::shouldSelectIndex)))
+ logicalAggregate(logicalProject(logicalOlapScan().when(this::shouldSelectIndex)))
  .then(agg -> {
  LogicalProject<LogicalOlapScan> project = agg.child();
  LogicalOlapScan scan = project.child();
@@ -134,7 +134,7 @@ public List<Rule> buildRules() {
  // filter could push down and project.
  // Aggregate(Project(Filter(Scan)))
  logicalAggregate(logicalProject(logicalFilter(logicalOlapScan()
- .when(LogicalOlapScan::shouldSelectIndex)))).then(agg -> {
+ .when(this::shouldSelectIndex)))).then(agg -> {
  LogicalProject<LogicalFilter<LogicalOlapScan>> project = agg.child();
  LogicalFilter<LogicalOlapScan> filter = project.child();
  LogicalOlapScan scan = filter.child();
@@ -157,7 +157,7 @@ public List<Rule> buildRules() {
  // filter can't push down
  // Aggregate(Filter(Project(Scan)))
  logicalAggregate(logicalFilter(logicalProject(logicalOlapScan()
- .when(LogicalOlapScan::shouldSelectIndex)))).then(agg -> {
+ .when(this::shouldSelectIndex)))).then(agg -> {
  LogicalFilter<LogicalProject<LogicalOlapScan>> filter = agg.child();
  LogicalProject<LogicalOlapScan> project = filter.child();
  LogicalOlapScan scan = project.child();
@@ -334,16 +334,18 @@ public PreAggStatus visitSum(Sum sum, CheckContext context) {
  return checkAggFunc(sum, AggregateType.SUM, extractSlotId(sum.child()), context, false);
  }
 
+ // TODO: select count(xxx) for duplicated-keys table.
  @Override
  public PreAggStatus visitCount(Count count, CheckContext context) {
- Optional<ExprId> exprIdOpt = extractSlotId(count.child());
- // Only count(distinct key_column) is supported.
- if (count.isDistinct() && exprIdOpt.isPresent() && context.exprIdToKeyColumn.containsKey(exprIdOpt.get())) {
- return PreAggStatus.on();
- } else {
- return PreAggStatus.off(String.format(
- "Count distinct is only valid for key columns, but meet %s.", count.toSql()));
+ // Now count(distinct key_column) is only supported for aggregate-keys and unique-keys OLAP table.
+ if (count.isDistinct()) {
+ Optional<ExprId> exprIdOpt = extractSlotId(count.child(0));
+ if (exprIdOpt.isPresent() && context.exprIdToKeyColumn.containsKey(exprIdOpt.get())) {
+ return PreAggStatus.on();
+ }
  }
+ return PreAggStatus.off(String.format(
+ "Count distinct is only valid for key columns, but meet %s.", count.toSql()));
  }
 
  private PreAggStatus checkAggFunc(

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/mv/SelectMaterializedIndexWithoutAggregate.java

@@ -55,7 +55,7 @@ public List<Rule> buildRules() {
  return ImmutableList.of(
  // project with pushdown filter.
  // Project(Filter(Scan))
- logicalProject(logicalFilter(logicalOlapScan().whenNot(LogicalOlapScan::isIndexSelected)))
+ logicalProject(logicalFilter(logicalOlapScan().when(this::shouldSelectIndex)))
  .then(project -> {
  LogicalFilter<LogicalOlapScan> filter = project.child();
  LogicalOlapScan scan = filter.child();
@@ -76,7 +76,7 @@ public List<Rule> buildRules() {
 
  // scan with filters could be pushdown.
  // Filter(Scan)
- logicalFilter(logicalOlapScan().whenNot(LogicalOlapScan::isIndexSelected))
+ logicalFilter(logicalOlapScan().when(this::shouldSelectIndex))
  .then(filter -> {
  LogicalOlapScan scan = filter.child();
  return filter.withChildren(select(scan, ImmutableSet::of, filter::getConjuncts));
@@ -85,7 +85,7 @@ public List<Rule> buildRules() {
 
  // project and scan.
  // Project(Scan)
- logicalProject(logicalOlapScan().whenNot(LogicalOlapScan::isIndexSelected))
+ logicalProject(logicalOlapScan().when(this::shouldSelectIndex))
  .then(project -> {
  LogicalOlapScan scan = project.child();
  return project.withChildren(
@@ -119,7 +119,7 @@ private LogicalOlapScan select(
  OlapTable table = scan.getTable();
  long baseIndexId = table.getBaseIndexId();
  int baseIndexKeySize = table.getKeyColumnsByIndexId(table.getBaseIndexId()).size();
- // No on aggregate on scan.
+ // No aggregate on scan.
  // So only base index and indexes that have all the keys could be used.
  List<MaterializedIndex> candidates = table.getVisibleIndex().stream()
  .filter(index -> index.getId() == baseIndexId

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/Count.java

@@ -20,7 +20,6 @@
 import org.apache.doris.nereids.exceptions.UnboundException;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
-import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.BigIntType;
 import org.apache.doris.nereids.types.DataType;
@@ -31,7 +30,7 @@
 import java.util.stream.Collectors;
 
 /** count agg function. */
-public class Count extends AggregateFunction implements UnaryExpression, AlwaysNotNullable {
+public class Count extends AggregateFunction implements AlwaysNotNullable {
 
  private final boolean isStar;
 

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java

@@ -22,7 +22,6 @@
 import org.apache.doris.catalog.Table;
 import org.apache.doris.nereids.memo.GroupExpression;
 import org.apache.doris.nereids.properties.LogicalProperties;
-import org.apache.doris.nereids.rules.mv.SelectMaterializedIndexWithAggregate;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.PreAggStatus;
@@ -74,7 +73,7 @@ public LogicalOlapScan(RelationId id, Table table, List<String> qualifier) {
  public LogicalOlapScan(RelationId id, Table table, List<String> qualifier,
  Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties,
  List<Long> selectedPartitionIdList, boolean partitionPruned, List<Long> candidateIndexIds,
- boolean rollupSelected, PreAggStatus preAggStatus) {
+ boolean indexSelected, PreAggStatus preAggStatus) {
  super(id, PlanType.LOGICAL_OLAP_SCAN, table, qualifier,
  groupExpression, logicalProperties, selectedPartitionIdList);
  // TODO: use CBO manner to select best index id, according to index's statistics info,
@@ -87,7 +86,7 @@ public LogicalOlapScan(RelationId id, Table table, List<String> qualifier,
  }
  this.partitionPruned = partitionPruned;
  this.candidateIndexIds = candidateIndexIds;
- this.indexSelected = rollupSelected;
+ this.indexSelected = indexSelected;
  this.preAggStatus = preAggStatus;
  }
 
@@ -172,20 +171,6 @@ public PreAggStatus getPreAggStatus() {
  return preAggStatus;
  }
 
- /**
- * Should apply {@link SelectMaterializedIndexWithAggregate} or not.
- */
- public boolean shouldSelectIndex() {
- switch (((OlapTable) table).getKeysType()) {
- case AGG_KEYS:
- case UNIQUE_KEYS:
- case DUP_KEYS:
- return !indexSelected;
- default:
- return false;
- }
- }
-
  @VisibleForTesting
  public Optional<String> getSelectedMaterializedIndexName() {
  return indexSelected ? Optional.ofNullable(((OlapTable) table).getIndexNameById(selectedIndexId))