Code review comment: use `thenApply(ctx -> {ctx.statementContext})` for better performance

Code before review:
public List<Rule> buildRules() {
return ImmutableList.of(
logicalJoin(any(), any())
.then(join -> {
if (hasEmptyRelationChild(join) && canReplaceJoinByEmptyRelation(join)
|| bothChildrenEmpty(join)) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
join.getOutput());
}
return null;
}
)
.toRule(RuleType.ELIMINATE_JOIN_ON_EMPTYRELATION),
logicalFilter(logicalEmptyRelation())
.then(filter -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
filter.getOutput())
).toRule(RuleType.ELIMINATE_FILTER_ON_EMPTYRELATION),
logicalAggregate(logicalEmptyRelation())
.when(agg -> !agg.getGroupByExpressions().isEmpty())
.then(agg -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
agg.getOutput())
).toRule(RuleType.ELIMINATE_AGG_ON_EMPTYRELATION),
logicalProject(logicalEmptyRelation())
.then(project ->
new LogicalEmptyRelation(ConnectContext.get().getStatementContext().getNextRelationId(),
project.getOutputs()
))
.toRule(RuleType.ELIMINATE_AGG_ON_EMPTYRELATION),
logicalUnion(multi()).then(union -> {
if (union.children().isEmpty()) {
return null;
}
ImmutableList.Builder<Plan> nonEmptyChildrenBuilder = ImmutableList.builder();
ImmutableList.Builder<List<SlotReference>> nonEmptyOutputsBuilder = ImmutableList.builder();
for (int i = 0; i < union.arity(); i++) {
if (!(union.child(i) instanceof EmptyRelation)) {
nonEmptyChildrenBuilder.add(union.child(i));
nonEmptyOutputsBuilder.add(union.getRegularChildOutput(i));
}
}
List<Plan> nonEmptyChildren = nonEmptyChildrenBuilder.build();
if (nonEmptyChildren.isEmpty()) {
if (union.getConstantExprsList().isEmpty()) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
union.getOutput());
} else {
return union.withChildrenAndTheirOutputs(ImmutableList.of(), ImmutableList.of());
}
} else if (nonEmptyChildren.size() == 1) {
if (union.getConstantExprsList().isEmpty()) {
Plan child = nonEmptyChildren.get(0);
List<Slot> unionOutput = union.getOutput();
List<Slot> childOutput = child.getOutput();
List<NamedExpression> projects = Lists.newArrayList();
for (int i = 0; i < unionOutput.size(); i++) {
ExprId id = unionOutput.get(i).getExprId();
Alias alias = new Alias(id, childOutput.get(i), unionOutput.get(i).getName());
projects.add(alias);
}
return new LogicalProject<>(projects, child);
}
}
if (union.children().size() != nonEmptyChildren.size()) {
return union.withChildrenAndTheirOutputs(nonEmptyChildren, nonEmptyOutputsBuilder.build());
} else {
return null;
}
}).toRule(RuleType.ELIMINATE_UNION_ON_EMPTYRELATION),
logicalTopN(logicalEmptyRelation())
.then(topn -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
topn.getOutput()))
.toRule(RuleType.ELIMINATE_TOPN_ON_EMPTYRELATION),
logicalSort(logicalEmptyRelation())
.then(sort -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
sort.getOutput()))
.toRule(RuleType.ELIMINATE_SORT_ON_EMPTYRELATION),
logicalIntersect(multi()).then(intersect -> {
List<Plan> emptyChildren = intersect.children().stream()
.filter(EmptyRelation.class::isInstance)
.collect(Collectors.toList());
if (emptyChildren.isEmpty()) {
return null;
} else {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
intersect.getOutput());
}
}).toRule(RuleType.ELIMINATE_INTERSECTION_ON_EMPTYRELATION),
logicalExcept(multi()).then(except -> {
Plan first = except.child(0);
if (first instanceof EmptyRelation) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
except.getOutput());
} else {
ImmutableList.Builder<Plan> nonEmptyChildrenBuilder = ImmutableList.builder();
ImmutableList.Builder<List<SlotReference>> nonEmptyOutputsBuilder = ImmutableList.builder();
for (int i = 0; i < except.arity(); i++) {
if (!(except.child(i) instanceof EmptyRelation)) {
nonEmptyChildrenBuilder.add(except.child(i));
nonEmptyOutputsBuilder.add(except.getRegularChildOutput(i));
}
}
List<Plan> nonEmptyChildren = nonEmptyChildrenBuilder.build();
if (nonEmptyChildren.size() == 1) {
Plan projectChild;
if (except.getQualifier() == SetOperation.Qualifier.DISTINCT) {
List<NamedExpression> firstOutputNamedExpressions = first.getOutput()
.stream().map(slot -> (NamedExpression) slot)
.collect(ImmutableList.toImmutableList());
projectChild = new LogicalAggregate<>(ImmutableList.copyOf(firstOutputNamedExpressions),
firstOutputNamedExpressions, true, Optional.empty(), first);
} else {
projectChild = first;
}
List<Slot> exceptOutput = except.getOutput();
List<Slot> projectInputSlots = projectChild.getOutput();
List<NamedExpression> projects = Lists.newArrayList();
for (int i = 0; i < exceptOutput.size(); i++) {
ExprId id = exceptOutput.get(i).getExprId();
Alias alias = new Alias(id, projectInputSlots.get(i), exceptOutput.get(i).getName());
projects.add(alias);
}
return new LogicalProject<>(projects, projectChild);
} else if (nonEmptyChildren.size() == except.children().size()) {
return null;
} else {
return except.withChildrenAndTheirOutputs(nonEmptyChildren, nonEmptyOutputsBuilder.build());
}
}
}).toRule(RuleType.ELIMINATE_EXCEPT_ON_EMPTYRELATION)
);
}

Line under review: .then(project ->

Code after addressing the review:
public List<Rule> buildRules() {
return ImmutableList.of(
logicalJoin(any(), any())
.when(join -> hasEmptyRelationChild(join) && canReplaceJoinByEmptyRelation(join)
|| bothChildrenEmpty(join))
.then(join -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
join.getOutput())
)
.toRule(RuleType.ELIMINATE_JOIN_ON_EMPTYRELATION),
logicalFilter(logicalEmptyRelation())
.then(filter -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
filter.getOutput())
).toRule(RuleType.ELIMINATE_FILTER_ON_EMPTYRELATION),
logicalAggregate(logicalEmptyRelation())
.when(agg -> !agg.getGroupByExpressions().isEmpty())
.then(agg -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
agg.getOutput())
).toRule(RuleType.ELIMINATE_AGG_ON_EMPTYRELATION),
logicalProject(logicalEmptyRelation())
.thenApply(ctx -> {
LogicalProject<? extends Plan> project = ctx.root;
return new LogicalEmptyRelation(ConnectContext.get().getStatementContext().getNextRelationId(),
project.getOutputs());
}).toRule(RuleType.ELIMINATE_AGG_ON_EMPTYRELATION),
logicalUnion(multi()).then(union -> {
if (union.children().isEmpty()) {
return null;
}
ImmutableList.Builder<Plan> nonEmptyChildrenBuilder = ImmutableList.builder();
ImmutableList.Builder<List<SlotReference>> nonEmptyOutputsBuilder = ImmutableList.builder();
for (int i = 0; i < union.arity(); i++) {
if (!(union.child(i) instanceof EmptyRelation)) {
nonEmptyChildrenBuilder.add(union.child(i));
nonEmptyOutputsBuilder.add(union.getRegularChildOutput(i));
}
}
List<Plan> nonEmptyChildren = nonEmptyChildrenBuilder.build();
if (nonEmptyChildren.isEmpty()) {
if (union.getConstantExprsList().isEmpty()) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
union.getOutput());
} else {
return union.withChildrenAndTheirOutputs(ImmutableList.of(), ImmutableList.of());
}
} else if (nonEmptyChildren.size() == 1) {
if (union.getConstantExprsList().isEmpty()) {
Plan child = nonEmptyChildren.get(0);
List<Slot> unionOutput = union.getOutput();
List<Slot> childOutput = child.getOutput();
List<NamedExpression> projects = Lists.newArrayList();
for (int i = 0; i < unionOutput.size(); i++) {
ExprId id = unionOutput.get(i).getExprId();
Alias alias = new Alias(id, childOutput.get(i), unionOutput.get(i).getName());
projects.add(alias);
}
return new LogicalProject<>(projects, child);
}
}
if (union.children().size() != nonEmptyChildren.size()) {
return union.withChildrenAndTheirOutputs(nonEmptyChildren, nonEmptyOutputsBuilder.build());
} else {
return null;
}
}).toRule(RuleType.ELIMINATE_UNION_ON_EMPTYRELATION),
logicalTopN(logicalEmptyRelation())
.then(topn -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
topn.getOutput()))
.toRule(RuleType.ELIMINATE_TOPN_ON_EMPTYRELATION),
logicalSort(logicalEmptyRelation())
.then(sort -> new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
sort.getOutput()))
.toRule(RuleType.ELIMINATE_SORT_ON_EMPTYRELATION),
logicalIntersect(multi()).then(intersect -> {
List<Plan> emptyChildren = intersect.children().stream()
.filter(EmptyRelation.class::isInstance)
.collect(Collectors.toList());
if (emptyChildren.isEmpty()) {
return null;
} else {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
intersect.getOutput());
}
}).toRule(RuleType.ELIMINATE_INTERSECTION_ON_EMPTYRELATION),
logicalLimit(logicalEmptyRelation())
.then(UnaryNode::child)
.toRule(RuleType.ELIMINATE_LIMIT_ON_EMPTY_RELATION),
logicalExcept(multi()).then(except -> {
Plan first = except.child(0);
if (first instanceof EmptyRelation) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
except.getOutput());
} else {
ImmutableList.Builder<Plan> nonEmptyChildrenBuilder = ImmutableList.builder();
ImmutableList.Builder<List<SlotReference>> nonEmptyOutputsBuilder = ImmutableList.builder();
for (int i = 0; i < except.arity(); i++) {
if (!(except.child(i) instanceof EmptyRelation)) {
nonEmptyChildrenBuilder.add(except.child(i));
nonEmptyOutputsBuilder.add(except.getRegularChildOutput(i));
}
}
List<Plan> nonEmptyChildren = nonEmptyChildrenBuilder.build();
if (nonEmptyChildren.size() == 1) {
Plan projectChild;
if (except.getQualifier() == SetOperation.Qualifier.DISTINCT) {
List<NamedExpression> firstOutputNamedExpressions = first.getOutput()
.stream().map(slot -> (NamedExpression) slot)
.collect(ImmutableList.toImmutableList());
projectChild = new LogicalAggregate<>(ImmutableList.copyOf(firstOutputNamedExpressions),
firstOutputNamedExpressions, true, Optional.empty(), first);
} else {
projectChild = first;
}
List<Slot> exceptOutput = except.getOutput();
List<Slot> projectInputSlots = projectChild.getOutput();
List<NamedExpression> projects = Lists.newArrayList();
for (int i = 0; i < exceptOutput.size(); i++) {
ExprId id = exceptOutput.get(i).getExprId();
Alias alias = new Alias(id, projectInputSlots.get(i), exceptOutput.get(i).getName());
projects.add(alias);
}
return new LogicalProject<>(projects, projectChild);
} else if (nonEmptyChildren.size() == except.children().size()) {
return null;
} else {
return except.withChildrenAndTheirOutputs(nonEmptyChildren, nonEmptyOutputsBuilder.build());
}
}
}).toRule(RuleType.ELIMINATE_EXCEPT_ON_EMPTYRELATION)
);
}