[Fix](variant) filter with variant access may lead to to partition/tablet prune fall through

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

@@ -240,11 +240,14 @@ public enum RuleType {
  MATERIALIZED_INDEX_FILTER_PROJECT_SCAN(RuleTypeClass.REWRITE),
 
  OLAP_SCAN_PARTITION_PRUNE(RuleTypeClass.REWRITE),
+
+ OLAP_SCAN_WITH_PROJECT_PARTITION_PRUNE(RuleTypeClass.REWRITE),
  FILE_SCAN_PARTITION_PRUNE(RuleTypeClass.REWRITE),
  PUSH_CONJUNCTS_INTO_JDBC_SCAN(RuleTypeClass.REWRITE),
  PUSH_CONJUNCTS_INTO_ODBC_SCAN(RuleTypeClass.REWRITE),
  PUSH_CONJUNCTS_INTO_ES_SCAN(RuleTypeClass.REWRITE),
  OLAP_SCAN_TABLET_PRUNE(RuleTypeClass.REWRITE),
+ OLAP_SCAN_WITH_PROJECT_TABLET_PRUNE(RuleTypeClass.REWRITE),
  PUSH_AGGREGATE_TO_OLAP_SCAN(RuleTypeClass.REWRITE),
  EXTRACT_SINGLE_TABLE_EXPRESSION_FROM_DISJUNCTION(RuleTypeClass.REWRITE),
  HIDE_ONE_ROW_RELATION_UNDER_UNION(RuleTypeClass.REWRITE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java

@@ -21,17 +21,21 @@
 import org.apache.doris.catalog.OlapTable;
 import org.apache.doris.catalog.PartitionInfo;
 import org.apache.doris.catalog.PartitionItem;
+import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.rules.Rule;
 import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.rules.expression.rules.PartitionPruner;
 import org.apache.doris.nereids.rules.expression.rules.PartitionPruner.PartitionTableType;
 import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
 import org.apache.doris.nereids.util.Utils;
 import org.apache.doris.qe.ConnectContext;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
 
 import java.util.ArrayList;
@@ -45,59 +49,75 @@
  * Used to prune partition of olap scan, should execute after SwapProjectAndFilter, MergeConsecutiveFilters,
  * MergeConsecutiveProjects and all predicate push down related rules.
  */
-public class PruneOlapScanPartition extends OneRewriteRuleFactory {
+public class PruneOlapScanPartition implements RewriteRuleFactory {
+ private <T extends Plan> Plan prunePartitions(CascadesContext ctx,
+ LogicalOlapScan scan, LogicalFilter<T> originalFilter) {
+ OlapTable table = scan.getTable();
+ Set<String> partitionColumnNameSet = Utils.execWithReturnVal(table::getPartitionColumnNames);
+ if (partitionColumnNameSet.isEmpty()) {
+ return originalFilter;
+ }
 
- @Override
- public Rule build() {
- return logicalFilter(logicalOlapScan()).when(p -> !p.child().isPartitionPruned()).thenApply(ctx -> {
- LogicalFilter<LogicalOlapScan> filter = ctx.root;
- LogicalOlapScan scan = filter.child();
- OlapTable table = scan.getTable();
- Set<String> partitionColumnNameSet = Utils.execWithReturnVal(table::getPartitionColumnNames);
- if (partitionColumnNameSet.isEmpty()) {
- return filter;
- }
+ List<Slot> output = scan.getOutput();
+ Map<String, Slot> scanOutput = Maps.newHashMapWithExpectedSize(output.size() * 2);
+ for (Slot slot : output) {
+ scanOutput.put(slot.getName().toLowerCase(), slot);
+ }
 
- List<Slot> output = scan.getOutput();
- Map<String, Slot> scanOutput = Maps.newHashMapWithExpectedSize(output.size() * 2);
- for (Slot slot : output) {
- scanOutput.put(slot.getName().toLowerCase(), slot);
+ PartitionInfo partitionInfo = table.getPartitionInfo();
+ List<Column> partitionColumns = partitionInfo.getPartitionColumns();
+ List<Slot> partitionSlots = new ArrayList<>(partitionColumns.size());
+ for (Column column : partitionColumns) {
+ Slot slot = scanOutput.get(column.getName().toLowerCase());
+ if (slot == null) {
+ return originalFilter;
+ } else {
+ partitionSlots.add(slot);
  }
+ }
 
- PartitionInfo partitionInfo = table.getPartitionInfo();
- List<Column> partitionColumns = partitionInfo.getPartitionColumns();
- List<Slot> partitionSlots = new ArrayList<>(partitionColumns.size());
- for (Column column : partitionColumns) {
- Slot slot = scanOutput.get(column.getName().toLowerCase());
- if (slot == null) {
- return filter;
- } else {
- partitionSlots.add(slot);
- }
- }
+ List<Long> manuallySpecifiedPartitions = scan.getManuallySpecifiedPartitions();
 
- List<Long> manuallySpecifiedPartitions = scan.getManuallySpecifiedPartitions();
+ Map<Long, PartitionItem> idToPartitions;
+ if (manuallySpecifiedPartitions.isEmpty()) {
+ idToPartitions = partitionInfo.getIdToItem(false);
+ } else {
+ Map<Long, PartitionItem> allPartitions = partitionInfo.getAllPartitions();
+ idToPartitions = allPartitions.keySet().stream()
+ .filter(id -> manuallySpecifiedPartitions.contains(id))
+ .collect(Collectors.toMap(Function.identity(), id -> allPartitions.get(id)));
+ }
+ List<Long> prunedPartitions = PartitionPruner.prune(
+ partitionSlots, originalFilter.getPredicate(), idToPartitions, ctx,
+ PartitionTableType.OLAP);
+ if (prunedPartitions.isEmpty()) {
+ return new LogicalEmptyRelation(
+ ConnectContext.get().getStatementContext().getNextRelationId(),
+ originalFilter.getOutput());
+ }
 
- Map<Long, PartitionItem> idToPartitions;
- if (manuallySpecifiedPartitions.isEmpty()) {
- idToPartitions = partitionInfo.getIdToItem(false);
- } else {
- Map<Long, PartitionItem> allPartitions = partitionInfo.getAllPartitions();
- idToPartitions = allPartitions.keySet().stream()
- .filter(id -> manuallySpecifiedPartitions.contains(id))
- .collect(Collectors.toMap(Function.identity(), id -> allPartitions.get(id)));
- }
- List<Long> prunedPartitions = PartitionPruner.prune(
- partitionSlots, filter.getPredicate(), idToPartitions, ctx.cascadesContext,
- PartitionTableType.OLAP);
+ LogicalOlapScan rewrittenScan = scan.withSelectedPartitionIds(prunedPartitions);
+ if (originalFilter.child() instanceof LogicalProject) {
+ LogicalProject<LogicalOlapScan> rewrittenProject
+ = (LogicalProject<LogicalOlapScan>) originalFilter.child()
+ .withChildren(ImmutableList.of(rewrittenScan));
+ return new LogicalFilter<>(originalFilter.getConjuncts(), rewrittenProject);
+ }
+ return originalFilter.withChildren(ImmutableList.of(rewrittenScan));
+ }
 
- if (prunedPartitions.isEmpty()) {
- return new LogicalEmptyRelation(
- ConnectContext.get().getStatementContext().getNextRelationId(),
- filter.getOutput());
- }
- LogicalOlapScan rewrittenScan = scan.withSelectedPartitionIds(prunedPartitions);
- return new LogicalFilter<>(filter.getConjuncts(), rewrittenScan);
- }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE);
+ @Override
+ public List<Rule> buildRules() {
+ return ImmutableList.of(
+ logicalFilter(logicalOlapScan()).when(p -> !p.child().isPartitionPruned()).thenApply(ctx -> {
+ return prunePartitions(ctx.cascadesContext, ctx.root.child(), ctx.root);
+ }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE),
+
+ logicalFilter(logicalProject(logicalOlapScan()))
+ .when(p -> !p.child().child().isPartitionPruned())
+ .when(p -> p.child().hasPushedDownToProjectionFunctions()).thenApply(ctx -> {
+ return prunePartitions(ctx.cascadesContext, ctx.root.child().child(), ctx.root);
+ }).toRule(RuleType.OLAP_SCAN_WITH_PROJECT_PARTITION_PRUNE)
+ );
  }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java

@@ -27,7 +27,10 @@
 import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionColumnFilterConverter;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
 import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.planner.HashDistributionPruner;
 import org.apache.doris.planner.PartitionColumnFilter;
@@ -46,34 +49,50 @@
 /**
  * prune bucket
  */
-public class PruneOlapScanTablet extends OneRewriteRuleFactory {
-
+public class PruneOlapScanTablet implements RewriteRuleFactory {
  @Override
- public Rule build() {
- return logicalFilter(logicalOlapScan())
- .then(filter -> {
- LogicalOlapScan olapScan = filter.child();
- OlapTable table = olapScan.getTable();
- Builder<Long> selectedTabletIdsBuilder = ImmutableList.builder();
- if (olapScan.getSelectedTabletIds().isEmpty()) {
- for (Long id : olapScan.getSelectedPartitionIds()) {
- Partition partition = table.getPartition(id);
- MaterializedIndex index = partition.getIndex(olapScan.getSelectedIndexId());
- selectedTabletIdsBuilder
- .addAll(getSelectedTabletIds(filter.getConjuncts(), index,
- olapScan.getSelectedIndexId() == olapScan.getTable()
- .getBaseIndexId(),
- partition.getDistributionInfo()));
- }
- } else {
- selectedTabletIdsBuilder.addAll(olapScan.getSelectedTabletIds());
- }
- List<Long> selectedTabletIds = selectedTabletIdsBuilder.build();
- if (new HashSet(selectedTabletIds).equals(new HashSet(olapScan.getSelectedTabletIds()))) {
- return null;
- }
- return filter.withChildren(olapScan.withSelectedTabletIds(selectedTabletIds));
- }).toRule(RuleType.OLAP_SCAN_TABLET_PRUNE);
+ public List<Rule> buildRules() {
+ return ImmutableList.of(
+ logicalFilter(logicalOlapScan())
+ .then(filter -> {
+ return pruneTablets(filter.child(), filter);
+ }).toRule(RuleType.OLAP_SCAN_TABLET_PRUNE),
+
+ logicalFilter(logicalProject(logicalOlapScan()))
+ .when(p -> p.child().hasPushedDownToProjectionFunctions()).then(filter -> {
+ return pruneTablets(filter.child().child(), filter);
+ }).toRule(RuleType.OLAP_SCAN_WITH_PROJECT_TABLET_PRUNE)
+ );
+ }
+
+ private <T extends Plan> Plan pruneTablets(LogicalOlapScan olapScan, LogicalFilter<T> originalFilter) {
+ OlapTable table = olapScan.getTable();
+ Builder<Long> selectedTabletIdsBuilder = ImmutableList.builder();
+ if (olapScan.getSelectedTabletIds().isEmpty()) {
+ for (Long id : olapScan.getSelectedPartitionIds()) {
+ Partition partition = table.getPartition(id);
+ MaterializedIndex index = partition.getIndex(olapScan.getSelectedIndexId());
+ selectedTabletIdsBuilder
+ .addAll(getSelectedTabletIds(originalFilter.getConjuncts(), index,
+ olapScan.getSelectedIndexId() == olapScan.getTable()
+ .getBaseIndexId(),
+ partition.getDistributionInfo()));
+ }
+ } else {
+ selectedTabletIdsBuilder.addAll(olapScan.getSelectedTabletIds());
+ }
+ List<Long> selectedTabletIds = selectedTabletIdsBuilder.build();
+ if (new HashSet(selectedTabletIds).equals(new HashSet(olapScan.getSelectedTabletIds()))) {
+ return null;
+ }
+ LogicalOlapScan rewrittenScan = olapScan.withSelectedTabletIds(selectedTabletIds);
+ if (originalFilter.child() instanceof LogicalProject) {
+ LogicalProject<LogicalOlapScan> rewrittenProject
+ = (LogicalProject<LogicalOlapScan>) originalFilter.child()
+ .withChildren(ImmutableList.of(rewrittenScan));
+ return new LogicalFilter<>(originalFilter.getConjuncts(), rewrittenProject);
+ }
+ return originalFilter.withChildren(rewrittenScan);
  }
 
  private Collection<Long> getSelectedTabletIds(Set<Expression> expressions,

regression-test/data/variant_p0/select_partition.out

@@ -0,0 +1,29 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !sql --
+1 {"a":1}
+
+-- !sql --
+7 {"a":2}
+
+-- !sql --
+1 {"a":1}
+7 {"a":2}
+
+-- !sql --
+7 {"a":2}
+
+-- !sql --
+16 {"a":3}
+
+-- !sql --
+16 {"a":3}
+
+-- !sql --
+16 {"a":3}
+
+-- !sql --
+1 {"a":1}
+
+-- !sql --
+6 {"x":"123"}
+