[opt](Nereids)choose better column for count star (#31227)

This commit is contained in:
minghong
2024-02-22 18:25:07 +08:00
committed by yiguolei
parent 9f8906269b
commit f1a7f9a70f
2 changed files with 79 additions and 3 deletions

View File

@ -80,8 +80,41 @@ import java.util.stream.IntStream;
* plan
*/
public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements CustomRewriter {
private Set<Slot> keys;
/**
* collect all columns used in expressions, which should not be pruned
*/
public static class KeyColumnCollector
extends DefaultPlanRewriter<JobContext> implements CustomRewriter {
public Set<Slot> keys = Sets.newHashSet();
@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
return plan.accept(this, jobContext);
}
@Override
public Plan visit(Plan plan, JobContext jobContext) {
for (Plan child : plan.children()) {
child.accept(this, jobContext);
}
plan.getExpressions().stream().filter(
expression -> !(expression instanceof SlotReference)
).forEach(
expression -> {
keys.addAll(expression.getInputSlots());
}
);
return plan;
}
}
@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
KeyColumnCollector keyColumnCollector = new KeyColumnCollector();
plan.accept(keyColumnCollector, jobContext);
keys = keyColumnCollector.keys;
return plan.accept(this, new PruneContext(plan.getOutputSet(), null));
}
@ -244,7 +277,7 @@ public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements
}
/** prune output */
public static <P extends Plan> P pruneOutput(P plan, List<NamedExpression> originOutput,
public <P extends Plan> P pruneOutput(P plan, List<NamedExpression> originOutput,
Function<List<NamedExpression>, P> withPrunedOutput, PruneContext context) {
if (originOutput.isEmpty()) {
return plan;
@ -254,7 +287,12 @@ public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements
.collect(ImmutableList.toImmutableList());
if (prunedOutputs.isEmpty()) {
NamedExpression minimumColumn = ExpressionUtils.selectMinimumColumn(originOutput);
List<NamedExpression> candidates = Lists.newArrayList(originOutput);
candidates.retainAll(keys);
if (candidates.isEmpty()) {
candidates = originOutput;
}
NamedExpression minimumColumn = ExpressionUtils.selectMinimumColumn(candidates);
prunedOutputs = ImmutableList.of(minimumColumn);
}