This is an automated email from the ASF dual-hosted git repository. jackietien pushed a commit to branch new_object_type in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 65ec0e2b3913f6a9740ed57869916cca584d168a Author: Weihao Li <[email protected]> AuthorDate: Mon Nov 3 16:51:52 2025 +0800 Add more optimizers for union (#16689) (cherry picked from commit 69cda2da662029b4f088ad5d7376d2f6902506a0) --- .../iterative/rule/PushLimitThroughUnion.java | 107 +++++++++++++++++++ .../iterative/rule/PushProjectionThroughUnion.java | 111 ++++++++++++++++++++ .../iterative/rule/PushTopKThroughUnion.java | 102 +++++++++++++++++++ .../iterative/rule/RemoveEmptyUnionBranches.java | 113 +++++++++++++++++++++ .../optimizations/LogicalOptimizeFactory.java | 24 +++-- ...{MergeUnionTest.java => UnionOptimizeTest.java} | 62 ++++++++++- 6 files changed, 511 insertions(+), 8 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushLimitThroughUnion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushLimitThroughUnion.java new file mode 100644 index 00000000000..0c2c11ac720 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushLimitThroughUnion.java @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule; + +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.Rule; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.LimitNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Captures; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Pattern; + +import com.google.common.collect.ImmutableList; + +import java.util.Optional; + +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.Limit.requiresPreSortedInputs; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.limit; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.source; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.union; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.QueryCardinalityUtil.isAtMost; +import static org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture.newCapture; + +/** + * Transforms: + * + * <pre> + * - Limit + * - Union + * - relation1 + * - relation2 + * .. + * </pre> + * + * Into: + * + * <pre> + * - Limit + * - Union + * - Limit + * - relation1 + * - Limit + * - relation2 + * .. + * </pre> + * + * Applies to LimitNode without ties only to avoid optimizer loop. + */ +public class PushLimitThroughUnion implements Rule<LimitNode> { + private static final Capture<UnionNode> CHILD = newCapture(); + + private static final Pattern<LimitNode> PATTERN = + limit() + .matching(limit -> !limit.isWithTies()) + .with(requiresPreSortedInputs().equalTo(false)) + .with(source().matching(union().capturedAs(CHILD))); + + @Override + public Pattern<LimitNode> getPattern() { + return PATTERN; + } + + @Override + public Result apply(LimitNode parent, Captures captures, Context context) { + UnionNode unionNode = captures.get(CHILD); + ImmutableList.Builder<PlanNode> builder = ImmutableList.builder(); + boolean shouldApply = false; + for (PlanNode child : unionNode.getChildren()) { + // This check is to ensure that we don't fire the optimizer if it was previously applied. + if (isAtMost(child, context.getLookup(), parent.getCount())) { + builder.add(child); + } else { + shouldApply = true; + builder.add( + new LimitNode( + context.getIdAllocator().genPlanNodeId(), + child, + parent.getCount(), + Optional.empty())); + } + } + + if (!shouldApply) { + return Result.empty(); + } + + return Result.ofPlanNode( + parent.replaceChildren(ImmutableList.of(unionNode.replaceChildren(builder.build())))); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushProjectionThroughUnion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushProjectionThroughUnion.java new file mode 100644 index 00000000000..4678ec16b48 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushProjectionThroughUnion.java @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule; + +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.Assignments; +import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.Rule; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ProjectNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Expression; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Captures; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Pattern; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableListMultimap; +import org.apache.tsfile.read.common.type.Type; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.apache.iotdb.db.queryengine.plan.relational.planner.ExpressionSymbolInliner.inlineSymbols; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.project; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.source; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.union; +import static org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture.newCapture; + +public class PushProjectionThroughUnion implements Rule<ProjectNode> { + private static final Capture<UnionNode> CHILD = newCapture(); + + private static final Pattern<ProjectNode> PATTERN = + project() + .matching(PushProjectionThroughUnion::nonTrivialProjection) + .with(source().matching(union().capturedAs(CHILD))); + + @Override + public Pattern<ProjectNode> getPattern() { + return PATTERN; + } + + @Override + public Result apply(ProjectNode parent, Captures captures, Context context) { + UnionNode child = captures.get(CHILD); + + // OutputLayout of the resultant Union, will be same as the layout of the Project + List<Symbol> outputLayout = parent.getOutputSymbols(); + + // Mapping from the output symbol to ordered list of symbols from each of the children + ImmutableListMultimap.Builder<Symbol, Symbol> mappings = ImmutableListMultimap.builder(); + + // sources for the resultant UnionNode + ImmutableList.Builder<PlanNode> outputSources = ImmutableList.builder(); + + for (int i = 0; i < child.getChildren().size(); i++) { + Map<Symbol, SymbolReference> outputToInput = + child.sourceSymbolMap(i); // Map: output of union -> input of this child to the union + Assignments.Builder assignments = + Assignments.builder(); // assignments for the new ProjectNode + + // mapping from current ProjectNode to new ProjectNode, used to identify the output layout + Map<Symbol, Symbol> projectSymbolMapping = new HashMap<>(); + + // Translate the assignments in the ProjectNode using symbols of the source of the UnionNode + for (Map.Entry<Symbol, Expression> entry : parent.getAssignments().entrySet()) { + Expression translatedExpression = inlineSymbols(outputToInput, entry.getValue()); + Type type = context.getSymbolAllocator().getTypes().getTableModelType(entry.getKey()); + Symbol symbol = context.getSymbolAllocator().newSymbol(translatedExpression, type); + assignments.put(symbol, translatedExpression); + projectSymbolMapping.put(entry.getKey(), symbol); + } + outputSources.add( + new ProjectNode( + context.getIdAllocator().genPlanNodeId(), + child.getChildren().get(i), + assignments.build())); + outputLayout.forEach(symbol -> mappings.put(symbol, projectSymbolMapping.get(symbol))); + } + + return Result.ofPlanNode( + new UnionNode( + parent.getPlanNodeId(), + outputSources.build(), + mappings.build(), + ImmutableList.copyOf(mappings.build().keySet()))); + } + + private static boolean nonTrivialProjection(ProjectNode project) { + return !project.getAssignments().getExpressions().stream() + .allMatch(SymbolReference.class::isInstance); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushTopKThroughUnion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushTopKThroughUnion.java new file mode 100644 index 00000000000..2064e7b6863 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/PushTopKThroughUnion.java @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule; + +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.Rule; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.SymbolMapper; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Captures; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Pattern; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; + +import java.util.Collections; +import java.util.Set; + +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Sets.intersection; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.source; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.topK; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.union; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.QueryCardinalityUtil.isAtMost; +import static org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Capture.newCapture; + +public class PushTopKThroughUnion implements Rule<TopKNode> { + private static final Capture<UnionNode> CHILD = newCapture(); + + private static final Pattern<TopKNode> PATTERN = + topK().with(source().matching(union().capturedAs(CHILD))); + + @Override + public Pattern<TopKNode> getPattern() { + return PATTERN; + } + + @Override + public Result apply(TopKNode topKNode, Captures captures, Context context) { + UnionNode unionNode = captures.get(CHILD); + + ImmutableList.Builder<PlanNode> children = ImmutableList.builder(); + + boolean shouldApply = false; + for (PlanNode child : unionNode.getChildren()) { + SymbolMapper.Builder symbolMapper = SymbolMapper.builder(); + Set<Symbol> sourceOutputSymbols = ImmutableSet.copyOf(child.getOutputSymbols()); + // This check is to ensure that we don't fire the optimizer if it was previously applied, + // which is the same as PushLimitThroughUnion. + if (isAtMost(child, context.getLookup(), topKNode.getCount())) { + children.add(child); + } else { + shouldApply = true; + for (Symbol unionOutput : unionNode.getOutputSymbols()) { + Set<Symbol> inputSymbols = + ImmutableSet.copyOf(unionNode.getSymbolMapping().get(unionOutput)); + Symbol unionInput = getLast(intersection(inputSymbols, sourceOutputSymbols)); + symbolMapper.put(unionOutput, unionInput); + } + children.add( + symbolMapper + .build() + .map( + topKNode, + Collections.singletonList(child), + context.getIdAllocator().genPlanNodeId())); + } + } + + if (!shouldApply) { + return Result.empty(); + } + + return Result.ofPlanNode( + topKNode.replaceChildren( + Collections.singletonList( + new UnionNode( + unionNode.getPlanNodeId(), + children.build(), + unionNode.getSymbolMapping(), + unionNode.getOutputSymbols())))); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/RemoveEmptyUnionBranches.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/RemoveEmptyUnionBranches.java new file mode 100644 index 00000000000..1b4021f7116 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/iterative/rule/RemoveEmptyUnionBranches.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule; + +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.Assignments; +import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.Rule; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ProjectNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Captures; +import org.apache.iotdb.db.queryengine.plan.relational.utils.matching.Pattern; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableListMultimap; +import com.google.common.collect.ListMultimap; + +import java.util.List; + +import static org.apache.iotdb.db.queryengine.plan.relational.planner.node.Patterns.union; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.QueryCardinalityUtil.isEmpty; + +/** + * Removes branches from a UnionNode that are guaranteed to produce 0 rows. + * + * <p>If there's only one branch left, it replaces the UnionNode with a projection to preserve the + * outputs of the union. + * + * <p>If all branches are empty, now we do the same process with one branch left case. + */ +public class RemoveEmptyUnionBranches implements Rule<UnionNode> { + private static final Pattern<UnionNode> PATTERN = union(); + + @Override + public Pattern<UnionNode> getPattern() { + return PATTERN; + } + + @Override + public Result apply(UnionNode node, Captures captures, Context context) { + int emptyBranches = 0; + ImmutableList.Builder<PlanNode> newChildrenBuilder = ImmutableList.builder(); + ImmutableListMultimap.Builder<Symbol, Symbol> outputsToInputsBuilder = + ImmutableListMultimap.builder(); + for (int i = 0; i < node.getChildren().size(); i++) { + PlanNode child = node.getChildren().get(i); + if (!isEmpty(child, context.getLookup())) { + newChildrenBuilder.add(child); + + for (Symbol column : node.getOutputSymbols()) { + outputsToInputsBuilder.put(column, node.getSymbolMapping().get(column).get(i)); + } + } else { + emptyBranches++; + } + } + + if (emptyBranches == 0) { + return Result.empty(); + } + + // restore after ValuesNode is introduced + /* + if (emptyBranches == node.getChildren().size()) { + return Result.ofPlanNode(new ValuesNode(node.getId(), node.getOutputSymbols(), ImmutableList.of())); + }*/ + + // Now we do the same process with one branch left case, choose the first child as preserved. + if (emptyBranches == node.getChildren().size()) { + Assignments.Builder assignments = Assignments.builder(); + for (Symbol column : node.getOutputSymbols()) { + assignments.put(column, node.getSymbolMapping().get(column).get(0).toSymbolReference()); + } + + return Result.ofPlanNode( + new ProjectNode(node.getPlanNodeId(), node.getChildren().get(0), assignments.build())); + } + + List<PlanNode> newChildren = newChildrenBuilder.build(); + ListMultimap<Symbol, Symbol> outputsToInputs = outputsToInputsBuilder.build(); + + if (newChildren.size() == 1) { + Assignments.Builder assignments = Assignments.builder(); + + outputsToInputs + .entries() + .forEach(entry -> assignments.put(entry.getKey(), entry.getValue().toSymbolReference())); + + return Result.ofPlanNode( + new ProjectNode(node.getPlanNodeId(), newChildren.get(0), assignments.build())); + } + + return Result.ofPlanNode( + new UnionNode(node.getPlanNodeId(), newChildren, outputsToInputs, node.getOutputSymbols())); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/LogicalOptimizeFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/LogicalOptimizeFactory.java index cdf99e55434..d70f67114b0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/LogicalOptimizeFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/LogicalOptimizeFactory.java @@ -68,7 +68,11 @@ import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.Pr import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PruneWindowColumns; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PushLimitThroughOffset; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PushLimitThroughProject; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PushLimitThroughUnion; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PushProjectionThroughUnion; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PushTopKThroughUnion; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.RemoveDuplicateConditions; +import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.RemoveEmptyUnionBranches; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.RemoveRedundantEnforceSingleRowNode; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.RemoveRedundantExists; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.RemoveRedundantIdentityProjections; @@ -141,8 +145,7 @@ public class LogicalOptimizeFactory { IterativeOptimizer columnPruningOptimizer = new IterativeOptimizer(plannerContext, ruleStats, columnPruningRules); - // Set<Rule<?>> projectionPushdownRules = ImmutableSet.of( - // new PushProjectionThroughUnion(), + Set<Rule<?>> projectionPushdownRules = ImmutableSet.of(new PushProjectionThroughUnion()); // new PushProjectionThroughExchange(), // // Dereference pushdown rules // new PushDownDereferencesThroughMarkDistinct(typeAnalyzer), @@ -181,7 +184,10 @@ public class LogicalOptimizeFactory { new IterativeOptimizer(plannerContext, ruleStats, simplifyOptimizerRules); Set<Rule<?>> limitPushdownRules = - ImmutableSet.of(new PushLimitThroughOffset(), new PushLimitThroughProject()); + ImmutableSet.of( + new PushLimitThroughOffset(), + new PushLimitThroughProject(), + new PushLimitThroughUnion()); ImmutableList.Builder<PlanOptimizer> optimizerBuilder = ImmutableList.builder(); @@ -199,12 +205,13 @@ public class LogicalOptimizeFactory { ruleStats, ImmutableSet.<Rule<?>>builder() .addAll(columnPruningRules) - // .addAll(projectionPushdownRules). + .addAll(projectionPushdownRules) // addAll(newUnwrapRowSubscript().rules()). // addAll(new PushCastIntoRow().rules()) .addAll( ImmutableSet.of( new ImplementTableFunctionSource(), + new RemoveEmptyUnionBranches(), new MergeFilters(), new InlineProjections(plannerContext), new RemoveRedundantIdentityProjections(), @@ -242,13 +249,13 @@ public class LogicalOptimizeFactory { plannerContext, ruleStats, ImmutableSet.<Rule<?>>builder() - // .addAll(projectionPushdownRules) + .addAll(projectionPushdownRules) .addAll(columnPruningRules) .addAll(limitPushdownRules) .addAll( ImmutableSet.of( new MergeUnion(), - // new RemoveEmptyUnionBranches(), + new RemoveEmptyUnionBranches(), new MergeFilters(), new RemoveTrivialFilters(), new MergeLimits(), @@ -331,7 +338,10 @@ public class LogicalOptimizeFactory { new IterativeOptimizer( plannerContext, ruleStats, - ImmutableSet.of(new MergeLimitWithSort(), new MergeLimitOverProjectWithSort())), + ImmutableSet.of( + new MergeLimitWithSort(), + new MergeLimitOverProjectWithSort(), + new PushTopKThroughUnion())), new ParallelizeGrouping()); this.planOptimizers = optimizerBuilder.build(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/MergeUnionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/UnionOptimizeTest.java similarity index 64% rename from iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/MergeUnionTest.java rename to iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/UnionOptimizeTest.java index d74e237082c..4fdbdb27b3d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/MergeUnionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/UnionOptimizeTest.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.plan.relational.analyzer; import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; +import org.apache.iotdb.db.queryengine.plan.relational.planner.PlanTester; import org.apache.iotdb.db.queryengine.plan.relational.planner.SymbolAllocator; import org.apache.iotdb.db.queryengine.plan.relational.planner.TableLogicalPlanner; @@ -31,11 +32,14 @@ import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.SESSION_INFO; import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.TEST_MATADATA; import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.limit; import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.output; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.project; import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.tableScan; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.topK; import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.union; -public class MergeUnionTest { +public class UnionOptimizeTest { @Test public void simpleLeftDeepMerge() { @@ -95,4 +99,60 @@ public class MergeUnionTest { tableScan("testdb.t3"), tableScan("testdb.t4"))))); } + + @Test + public void pushLimitThroughUnionTest() { + assertPlan( + new PlanTester() + .createPlan("(select tag1 from t1) union all (select tag1 from t2) limit 1"), + output( + limit(1, (union(limit(1, tableScan("testdb.t1")), limit(1, tableScan("testdb.t2"))))))); + } + + @Test + public void pushProjectionThroughUnionTest() { + assertPlan( + new PlanTester() + .createPlan("select s1 + 1 from ((select s1 from t1) union all (select s1 from t2)) "), + output((union(project(tableScan("testdb.t1")), project(tableScan("testdb.t2")))))); + } + + @Test + public void pushTopKThroughUnionTest() { + assertPlan( + new PlanTester() + .createPlan( + "(select tag1 from t1) union all (select tag1 from t2) order by tag1 limit 10"), + output(topK((union(topK(tableScan("testdb.t1")), topK(tableScan("testdb.t2"))))))); + } + + @Test + public void removeEmptyUnionBranchesTest1() { + // Normal case + assertPlan( + new PlanTester() + .createPlan( + "(select tag1 from t1 limit 0) union all (select tag1 from t2) union all (select tag1 from t3)"), + output((union(tableScan("testdb.t2"), tableScan("testdb.t3"))))); + } + + @Test + public void removeEmptyUnionBranchesTest2() { + // One non-empty branch + assertPlan( + new PlanTester() + .createPlan( + "(select tag1 from t1 limit 0) union all (select tag1 from t2 limit 0) union all (select tag1 from t3 limit 1)"), + output(limit(1, tableScan("testdb.t3")))); + } + + @Test + public void removeEmptyUnionBranchesTest3() { + // All branches are empty + assertPlan( + new PlanTester() + .createPlan( + "(select tag1 from t1 limit 0) union all (select tag1 from t2 limit 0) union all (select tag1 from t3 limit 0)"), + output(limit(0, tableScan("testdb.t1")))); + } }
