diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index 80e524e11f0e..e6850f26f9a7 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -73,6 +73,7 @@ private PinotQueryRuleSets() { // join and semi-join rules CoreRules.PROJECT_TO_SEMI_JOIN, + PinotSeminJoinDistinctProjectRule.INSTANCE, // convert non-all union into all-union + distinct CoreRules.UNION_TO_DISTINCT, diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java new file mode 100644 index 000000000000..f724e6d7a10d --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java @@ -0,0 +1,68 @@ +/** + * 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.pinot.calcite.rel.rules; + +import java.util.List; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.calcite.rel.logical.LogicalJoin; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.tools.RelBuilder; +import org.apache.calcite.tools.RelBuilderFactory; + + +/** + * Special rule for Pinot, this rule always append a distinct to the + * {@link org.apache.calcite.rel.logical.LogicalProject} on top of a Semi join + * {@link org.apache.calcite.rel.core.Join} to ensure the correctness of the query. + */ +public class PinotSeminJoinDistinctProjectRule extends RelOptRule { + public static final PinotSeminJoinDistinctProjectRule INSTANCE = + new PinotSeminJoinDistinctProjectRule(PinotRuleUtils.PINOT_REL_FACTORY); + + public PinotSeminJoinDistinctProjectRule(RelBuilderFactory factory) { + super(operand(LogicalJoin.class, some(operand(LogicalProject.class, any()), operand(LogicalProject.class, any()))), + factory, null); + } + + @Override + public void onMatch(RelOptRuleCall call) { + LogicalJoin join = call.rel(0); + if (join.getJoinType() != JoinRelType.SEMI) { + return; + } + LogicalProject leftProject = call.rel(1); + LogicalProject rightProject = call.rel(2); + + if (rightProject.getProjects().size() != 1) { + return; + } + RelNode newRightProject = insertDistinctToProject(call, rightProject); + call.transformTo(join.copy(join.getTraitSet(), List.of(leftProject, newRightProject))); + } + + private RelNode insertDistinctToProject(RelOptRuleCall call, LogicalProject project) { + RelBuilder relBuilder = call.builder(); + relBuilder.push(project); + relBuilder.distinct(); + return relBuilder.build(); + } +}