database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature](Nereids) elimite inner join by foreign key (#28486)

Browse Source
This commit is contained in:
谢健
2023-12-21 12:58:55 +08:00
committed by GitHub
parent db523dafcb
commit 223454d1db
12 changed files with 678 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
fe/fe-core/src/main/java/org/apache/doris/catalog/TableIf.java
View File

@ -165,6 +165,42 @@ public interface TableIf {
throw new RuntimeException(String.format("Not implemented constraint for table %s", this));
}
default Set<ForeignKeyConstraint> getForeignKeyConstraints() {
readLock();
try {
return getConstraintsMap().values().stream()
.filter(ForeignKeyConstraint.class::isInstance)
.map(ForeignKeyConstraint.class::cast)
.collect(ImmutableSet.toImmutableSet());
} finally {
readUnlock();
}
}
default Set<PrimaryKeyConstraint> getPrimaryKeyConstraints() {
readLock();
try {
return getConstraintsMap().values().stream()
.filter(PrimaryKeyConstraint.class::isInstance)
.map(PrimaryKeyConstraint.class::cast)
.collect(ImmutableSet.toImmutableSet());
} finally {
readUnlock();
}
}
default Set<UniqueConstraint> getUniqueConstraints() {
readLock();
try {
return getConstraintsMap().values().stream()
.filter(UniqueConstraint.class::isInstance)
.map(UniqueConstraint.class::cast)
.collect(ImmutableSet.toImmutableSet());
} finally {
readUnlock();
}
}
// Note this function is not thread safe
default void checkConstraintNotExistence(String name, Constraint primaryKeyConstraint,
Map<String, Constraint> constraintMap) {

9
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ForeignKeyConstraint.java
View File

@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableMap.Builder;
import com.google.common.collect.ImmutableSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
public class ForeignKeyConstraint extends Constraint {
@ -61,6 +62,14 @@ public class ForeignKeyConstraint extends Constraint {
return foreignToReference.get(column);
}
public Map<Column, Column> getForeignToPrimary(TableIf curTable) {
ImmutableMap.Builder<Column, Column> columnBuilder = new ImmutableMap.Builder<>();
TableIf refTable = referencedTable.toTableIf();
foreignToReference.forEach((k, v) ->
columnBuilder.put(curTable.getColumn(k), refTable.getColumn(v)));
return columnBuilder.build();
}
public Column getReferencedColumn(String column) {
return getReferencedTable().getColumn(getReferencedColumnName(column));
}

5
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/PrimaryKeyConstraint.java
View File

@ -17,6 +17,7 @@
package org.apache.doris.catalog.constraint;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.TableIf;
import com.google.common.base.Objects;
@ -42,6 +43,10 @@ public class PrimaryKeyConstraint extends Constraint {
return columns;
}
public ImmutableSet<Column> getPrimaryKeys(TableIf table) {
return columns.stream().map(table::getColumn).collect(ImmutableSet.toImmutableSet());
}
public void addForeignTable(TableIf table) {
foreignTables.add(new TableIdentifier(table));
}

7
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/UniqueConstraint.java
View File

@ -17,6 +17,9 @@
package org.apache.doris.catalog.constraint;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.TableIf;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableSet;
@ -34,6 +37,10 @@ public class UniqueConstraint extends Constraint {
return columns;
}
public ImmutableSet<Column> getUniqueKeys(TableIf table) {
return columns.stream().map(table::getColumn).collect(ImmutableSet.toImmutableSet());
}
@Override
public boolean equals(Object o) {
if (this == o) {

3
fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java
View File

@ -55,6 +55,7 @@ import org.apache.doris.nereids.rules.rewrite.EliminateAssertNumRows;
import org.apache.doris.nereids.rules.rewrite.EliminateDedupJoinCondition;
import org.apache.doris.nereids.rules.rewrite.EliminateEmptyRelation;
import org.apache.doris.nereids.rules.rewrite.EliminateFilter;
import org.apache.doris.nereids.rules.rewrite.EliminateJoinByFK;
import org.apache.doris.nereids.rules.rewrite.EliminateJoinCondition;
import org.apache.doris.nereids.rules.rewrite.EliminateLimit;
import org.apache.doris.nereids.rules.rewrite.EliminateNotNull;
@ -285,6 +286,8 @@ public class Rewriter extends AbstractBatchJobExecutor {
custom(RuleType.PUSH_DOWN_DISTINCT_THROUGH_JOIN, PushDownDistinctThroughJoin::new)
),
// this rule should invoke after infer predicate and push down distinct, and before push down limit
custom(RuleType.ELIMINATE_JOIN_BY_FOREIGN_KEY, EliminateJoinByFK::new),
topic("Limit optimization",
// TODO: the logical plan should not contains any phase information,
// we should refactor like AggregateStrategies, e.g. LimitStrategies,

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

@ -200,6 +200,7 @@ public enum RuleType {
ELIMINATE_LIMIT_ON_EMPTY_RELATION(RuleTypeClass.REWRITE),
ELIMINATE_FILTER(RuleTypeClass.REWRITE),
ELIMINATE_JOIN(RuleTypeClass.REWRITE),
ELIMINATE_JOIN_BY_FOREIGN_KEY(RuleTypeClass.REWRITE),
ELIMINATE_JOIN_CONDITION(RuleTypeClass.REWRITE),
ELIMINATE_FILTER_ON_ONE_RELATION(RuleTypeClass.REWRITE),
ELIMINATE_SEMI_JOIN(RuleTypeClass.REWRITE),

333
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateJoinByFK.java Normal file
View File

@ -0,0 +1,333 @@
// 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.doris.nereids.rules.rewrite;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.Table;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.nereids.jobs.JobContext;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.IsNull;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Not;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
import org.apache.doris.nereids.trees.plans.logical.LogicalLimit;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalSort;
import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
import org.apache.doris.nereids.trees.plans.logical.LogicalWindow;
import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter;
import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
import org.apache.doris.nereids.util.ImmutableEquivalenceSet;
import com.google.common.collect.BiMap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMap.Builder;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
/**
* Eliminate join by foreign.
*/
public class EliminateJoinByFK extends DefaultPlanRewriter<JobContext> implements CustomRewriter {
@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
EliminateJoinByFKHelper helper = new EliminateJoinByFKHelper();
return helper.rewriteRoot(plan, jobContext);
}
private static class EliminateJoinByFKHelper
extends DefaultPlanRewriter<ForeignKeyContext> implements CustomRewriter {
@Override
public Plan rewriteRoot(Plan plan, JobContext jobContext) {
return plan.accept(this, new ForeignKeyContext());
}
@Override
public Plan visit(Plan plan, ForeignKeyContext context) {
Plan newPlan = visitChildren(this, plan, context);
context.passThroughPlan(plan);
return newPlan;
}
@Override
public Plan visitLogicalRelation(LogicalRelation relation, ForeignKeyContext context) {
if (!(relation instanceof LogicalCatalogRelation)) {
return relation;
}
if (!(((LogicalCatalogRelation) relation).getTable() instanceof Table)) {
return relation;
}
context.putAllForeignKeys(((LogicalCatalogRelation) relation).getTable());
relation.getOutput().stream()
.filter(SlotReference.class::isInstance)
.map(SlotReference.class::cast)
.forEach(context::putSlot);
return relation;
}
private boolean canEliminate(LogicalJoin<?, ?> join, BiMap<Slot, Slot> equalSlots, ForeignKeyContext context) {
if (!join.getOtherJoinConjuncts().isEmpty()) {
return false;
}
if (!join.getJoinType().isInnerJoin() && !join.getJoinType().isSemiJoin()) {
return false;
}
return context.satisfyConstraint(equalSlots, join);
}
private boolean isForeignKeyAndUnique(Plan plan,
Set<Slot> keySet, ForeignKeyContext context) {
boolean unique = keySet.stream()
.allMatch(s -> plan.getLogicalProperties().getFunctionalDependencies().isUnique(s));
return unique && context.isForeignKey(keySet);
}
private @Nullable Map<Expression, Expression> tryGetOutputToChildMap(Plan child,
Set<Slot> output, BiMap<Slot, Slot> equalSlots) {
Set<Slot> residual = Sets.difference(output, child.getOutputSet());
if (equalSlots.keySet().containsAll(residual)) {
return residual.stream()
.collect(ImmutableMap.toImmutableMap(e -> e, equalSlots::get));
}
if (equalSlots.values().containsAll(residual)) {
return residual.stream()
.collect(ImmutableMap.toImmutableMap(e -> e, e -> equalSlots.inverse().get(e)));
}
return null;
}
private Plan applyNullCompensationFilter(Plan child, Set<Slot> childSlots) {
Set<Expression> predicates = childSlots.stream()
.filter(ExpressionTrait::nullable)
.map(s -> new Not(new IsNull(s)))
.collect(ImmutableSet.toImmutableSet());
if (predicates.isEmpty()) {
return child;
}
return new LogicalFilter<>(predicates, child);
}
private @Nullable Plan tryConstructPlanWithJoinChild(LogicalProject<LogicalJoin<?, ?>> project, Plan child,
BiMap<Slot, Slot> equalSlots, ForeignKeyContext context) {
Set<Slot> output = project.getInputSlots();
Set<Slot> keySet = child.getOutputSet().containsAll(equalSlots.keySet())
? equalSlots.keySet()
: equalSlots.values();
Map<Expression, Expression> outputToRight = tryGetOutputToChildMap(child, output, equalSlots);
if (outputToRight != null && isForeignKeyAndUnique(child, keySet, context)) {
List<NamedExpression> newProjects = project.getProjects().stream()
.map(e -> outputToRight.containsKey(e)
? new Alias(e.getExprId(), outputToRight.get(e), e.toSql())
: (NamedExpression) e.rewriteUp(s -> outputToRight.getOrDefault(s, s)))
.collect(ImmutableList.toImmutableList());
return project.withProjects(newProjects)
.withChildren(applyNullCompensationFilter(child, keySet));
}
return null;
}
// Right now we only support eliminate inner join, which should meet the following condition:
// 1. only contain null-reject equal condition, and which all meet fk-pk constraint
// 2. only output foreign table output or can be converted to foreign table output
// 3. foreign key is unique
// 4. if foreign key is null, add a isNotNull predicate for null-reject join condition
private Plan eliminateJoin(LogicalProject<LogicalJoin<?, ?>> project, ForeignKeyContext context) {
LogicalJoin<?, ?> join = project.child();
ImmutableEquivalenceSet<Slot> equalSet = join.getEqualSlots();
BiMap<Slot, Slot> equalSlots = equalSet.tryToMap();
if (equalSlots == null) {
return project;
}
if (!canEliminate(join, equalSlots, context)) {
return project;
}
Plan keepLeft = tryConstructPlanWithJoinChild(project, join.left(), equalSlots, context);
if (keepLeft != null) {
return keepLeft;
}
Plan keepRight = tryConstructPlanWithJoinChild(project, join.right(), equalSlots, context);
if (keepRight != null) {
return keepRight;
}
return project;
}
@Override
public Plan visitLogicalProject(LogicalProject<?> project, ForeignKeyContext context) {
project = visitChildren(this, project, context);
for (NamedExpression expression : project.getProjects()) {
if (expression instanceof Alias && expression.child(0) instanceof Slot) {
context.putAlias(expression.toSlot(), (Slot) expression.child(0));
}
}
if (project.child() instanceof LogicalJoin<?, ?>) {
return eliminateJoin((LogicalProject<LogicalJoin<?, ?>>) project, context);
}
return project;
}
}
private static class ForeignKeyContext {
static Set<Class<?>> propagatePrimaryKeyOperator = ImmutableSet
.<Class<?>>builder()
.add(LogicalProject.class)
.add(LogicalSort.class)
.add(LogicalJoin.class)
.build();
static Set<Class<?>> propagateForeignKeyOperator = ImmutableSet
.<Class<?>>builder()
.add(LogicalProject.class)
.add(LogicalSort.class)
.add(LogicalJoin.class)
.add(LogicalFilter.class)
.add(LogicalTopN.class)
.add(LogicalLimit.class)
.add(LogicalAggregate.class)
.add(LogicalWindow.class)
.build();
Set<Map<Column, Column>> constraints = new HashSet<>();
Set<Column> foreignKeys = new HashSet<>();
Set<Column> primaryKeys = new HashSet<>();
Map<Slot, Column> slotToColumn = new HashMap<>();
Map<Column, Set<LogicalJoin<?, ?>>> columnWithJoin = new HashMap<>();
Map<Column, Set<Expression>> columnWithPredicates = new HashMap<>();
public void putAllForeignKeys(TableIf table) {
table.getForeignKeyConstraints().forEach(c -> {
Map<Column, Column> constraint = c.getForeignToPrimary(table);
constraints.add(c.getForeignToPrimary(table));
foreignKeys.addAll(constraint.keySet());
primaryKeys.addAll(constraint.values());
});
}
public boolean isForeignKey(Set<Slot> key) {
return foreignKeys.containsAll(
key.stream().map(s -> slotToColumn.get(s)).collect(Collectors.toSet()));
}
public void putSlot(SlotReference slot) {
if (!slot.getColumn().isPresent()) {
return;
}
Column c = slot.getColumn().get();
slotToColumn.put(slot, c);
}
public void putAlias(Slot newSlot, Slot originSlot) {
if (slotToColumn.containsKey(originSlot)) {
slotToColumn.put(newSlot, slotToColumn.get(originSlot));
}
}
public void passThroughPlan(Plan plan) {
Set<Column> output = plan.getOutput().stream()
.filter(slotToColumn::containsKey)
.map(s -> slotToColumn.get(s))
.collect(ImmutableSet.toImmutableSet());
if (plan instanceof LogicalJoin) {
output.forEach(c ->
columnWithJoin.computeIfAbsent(c, v -> Sets.newHashSet((LogicalJoin<?, ?>) plan)));
return;
}
if (plan instanceof LogicalFilter) {
output.forEach(c -> {
columnWithPredicates.computeIfAbsent(c, v -> new HashSet<>());
columnWithPredicates.get(c).addAll(((LogicalFilter<?>) plan).getConjuncts());
});
return;
}
if (!propagatePrimaryKeyOperator.contains(plan.getClass())) {
output.forEach(primaryKeys::remove);
}
if (!propagateForeignKeyOperator.contains(plan.getClass())) {
output.forEach(foreignKeys::remove);
}
}
public boolean satisfyConstraint(BiMap<Slot, Slot> equalSlots, LogicalJoin<?, ?> join) {
ImmutableMap.Builder<Column, Column> foreignToPrimaryBuilder = new Builder<>();
for (Entry<Slot, Slot> entry : equalSlots.entrySet()) {
Slot left = entry.getKey();
Slot right = entry.getValue();
if (!slotToColumn.containsKey(left) || !slotToColumn.containsKey(right)) {
return false;
}
Column leftColumn = slotToColumn.get(left);
Column rightColumn = slotToColumn.get(right);
if (foreignKeys.contains(leftColumn)) {
foreignToPrimaryBuilder.put(leftColumn, rightColumn);
} else if (foreignKeys.contains(rightColumn)) {
foreignToPrimaryBuilder.put(rightColumn, leftColumn);
} else {
return false;
}
}
Map<Column, Column> foreignToPrimary = foreignToPrimaryBuilder.build();
// The primary key can only contain join that may be eliminated
if (!foreignToPrimary.values().stream().allMatch(p ->
columnWithJoin.get(p).size() == 1 && columnWithJoin.get(p).iterator().next() == join)) {
return false;
}
// The foreign key's filters must contain primary filters
if (!isPredicateCompatible(equalSlots, foreignToPrimary)) {
return false;
}
return constraints.contains(foreignToPrimary);
}
private boolean isPredicateCompatible(BiMap<Slot, Slot> equalSlots, Map<Column, Column> foreignToPrimary) {
return foreignToPrimary.entrySet().stream().allMatch(fp -> {
BiMap<Slot, Slot> primarySlotToForeign = equalSlots.keySet().stream()
.map(slotToColumn::get)
.anyMatch(primaryKeys::contains)
? equalSlots :
equalSlots.inverse();
if (!columnWithPredicates.containsKey(fp.getValue())) {
return true;
}
Set<Expression> primaryPredicates = columnWithPredicates.get(fp.getValue()).stream()
.map(e -> e.rewriteUp(
s -> s instanceof Slot ? primarySlotToForeign.getOrDefault(s, (Slot) s) : s))
.collect(Collectors.toSet());
return columnWithPredicates.get(fp.getKey()).containsAll(primaryPredicates);
});
}
}
}

25
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCatalogRelation.java
View File

@ -17,6 +17,7 @@
package org.apache.doris.nereids.trees.plans.logical;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.DatabaseIf;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.OlapTable;
@ -42,6 +43,7 @@ import org.apache.commons.lang3.StringUtils;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;
/**
@ -122,8 +124,9 @@ public abstract class LogicalCatalogRelation extends LogicalRelation implements
@Override
public FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplier) {
Builder fdBuilder = new Builder();
Set<Slot> output = ImmutableSet.copyOf(outputSupplier.get());
if (table instanceof OlapTable && ((OlapTable) table).getKeysType().isAggregationFamily()) {
ImmutableSet<Slot> slotSet = computeOutput().stream()
ImmutableSet<Slot> slotSet = output.stream()
.filter(SlotReference.class::isInstance)
.map(SlotReference.class::cast)
.filter(s -> s.getColumn().isPresent()
@ -131,6 +134,26 @@ public abstract class LogicalCatalogRelation extends LogicalRelation implements
.collect(ImmutableSet.toImmutableSet());
fdBuilder.addUniqueSlot(slotSet);
}
table.getPrimaryKeyConstraints().forEach(c -> {
Set<Column> columns = c.getPrimaryKeys(this.getTable());
ImmutableSet<Slot> slotSet = output.stream()
.filter(SlotReference.class::isInstance)
.map(SlotReference.class::cast)
.filter(s -> s.getColumn().isPresent()
&& columns.contains(s.getColumn().get()))
.collect(ImmutableSet.toImmutableSet());
fdBuilder.addUniqueSlot(slotSet);
});
table.getUniqueConstraints().forEach(c -> {
Set<Column> columns = c.getUniqueKeys(this.getTable());
ImmutableSet<Slot> slotSet = output.stream()
.filter(SlotReference.class::isInstance)
.map(SlotReference.class::cast)
.filter(s -> s.getColumn().isPresent()
&& columns.contains(s.getColumn().get()))
.collect(ImmutableSet.toImmutableSet());
fdBuilder.addUniqueSlot(slotSet);
});
return fdBuilder.build();
}
}

20
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalJoin.java
View File

@ -24,6 +24,7 @@ import org.apache.doris.nereids.properties.FunctionalDependencies;
import org.apache.doris.nereids.properties.FunctionalDependencies.Builder;
import org.apache.doris.nereids.properties.LogicalProperties;
import org.apache.doris.nereids.rules.exploration.join.JoinReorderContext;
import org.apache.doris.nereids.trees.expressions.EqualPredicate;
import org.apache.doris.nereids.trees.expressions.EqualTo;
import org.apache.doris.nereids.trees.expressions.ExprId;
import org.apache.doris.nereids.trees.expressions.Expression;
@ -36,6 +37,7 @@ import org.apache.doris.nereids.trees.plans.PlanType;
import org.apache.doris.nereids.trees.plans.algebra.Join;
import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.ImmutableEquivalenceSet;
import org.apache.doris.nereids.util.JoinUtils;
import org.apache.doris.nereids.util.Utils;
@ -454,6 +456,24 @@ public class LogicalJoin<LEFT_CHILD_TYPE extends Plan, RIGHT_CHILD_TYPE extends
return fdBuilder.build();
}
/**
* get Equal slot from join
*/
public ImmutableEquivalenceSet<Slot> getEqualSlots() {
// TODO: Use fd in the future
if (!joinType.isInnerJoin() && !joinType.isSemiJoin()) {
return ImmutableEquivalenceSet.of();
}
ImmutableEquivalenceSet.Builder<Slot> builder = new ImmutableEquivalenceSet.Builder<>();
hashJoinConjuncts.stream()
.filter(e -> e instanceof EqualPredicate
&& e.child(0) instanceof Slot
&& e.child(1) instanceof Slot)
.forEach(e ->
builder.addEqualPair((Slot) e.child(0), (Slot) e.child(1)));
return builder.build();
}
@Override
public JSONObject toJson() {
JSONObject logicalJoin = super.toJson();

101
fe/fe-core/src/main/java/org/apache/doris/nereids/util/ImmutableEquivalenceSet.java Normal file
View File

@ -0,0 +1,101 @@
// 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.doris.nereids.util;
import com.google.common.collect.BiMap;
import com.google.common.collect.ImmutableBiMap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import javax.annotation.Nullable;
/**
* EquivalenceSet
*/
public class ImmutableEquivalenceSet<T> {
final Map<T, T> root;
ImmutableEquivalenceSet(Map<T, T> root) {
this.root = ImmutableMap.copyOf(root);
}
public static <T> ImmutableEquivalenceSet<T> of() {
return new ImmutableEquivalenceSet<>(ImmutableMap.of());
}
/**
* Builder of ImmutableEquivalenceSet
*/
public static class Builder<T> {
final Map<T, T> parent = new HashMap<>();
public void addEqualPair(T a, T b) {
parent.computeIfAbsent(b, v -> v);
parent.computeIfAbsent(a, v -> v);
union(a, b);
}
private void union(T a, T b) {
T root1 = findRoot(a);
T root2 = findRoot(b);
if (root1 != root2) {
parent.put(b, root1);
findRoot(b);
}
}
private T findRoot(T a) {
if (!parent.get(a).equals(a)) {
parent.put(a, findRoot(parent.get(a)));
}
return parent.get(a);
}
public ImmutableEquivalenceSet<T> build() {
parent.keySet().forEach(this::findRoot);
return new ImmutableEquivalenceSet<>(parent);
}
}
/**
* cal equal set for a
*/
public Set<T> calEqualSet(T a) {
T ra = root.get(a);
return root.keySet().stream()
.filter(t -> root.get(t).equals(ra))
.collect(ImmutableSet.toImmutableSet());
}
/**
* try to convert it to a map, such as a = b c = d.
* When meets a = b a = c, return null
*/
public @Nullable BiMap<T, T> tryToMap() {
if (root.values().stream().distinct().count() * 2 != root.size()) {
return null;
}
return root.keySet().stream()
.filter(t -> !root.get(t).equals(t))
.collect(ImmutableBiMap.toImmutableBiMap(t -> t, root::get));
}
}