CheckAfterRewrite.java
// 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.analysis;
import org.apache.doris.catalog.Type;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.ComparisonPredicate;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.Match;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotNotFromChildren;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.expressions.WindowExpression;
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
import org.apache.doris.nereids.trees.expressions.functions.generator.TableGeneratingFunction;
import org.apache.doris.nereids.trees.expressions.functions.scalar.GroupingScalarFunction;
import org.apache.doris.nereids.trees.expressions.functions.window.WindowFunction;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.algebra.Generate;
import org.apache.doris.nereids.trees.plans.algebra.SetOperation;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
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.LogicalOlapScan;
import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat;
import org.apache.doris.nereids.trees.plans.logical.LogicalSetOperation;
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.LogicalUnion;
import org.apache.doris.nereids.trees.plans.logical.LogicalWindow;
import org.apache.doris.nereids.types.DataType;
import org.apache.doris.nereids.types.VariantType;
import org.apache.doris.qe.ConnectContext;
import com.google.common.collect.Sets;
import org.apache.commons.lang3.StringUtils;
import org.roaringbitmap.RoaringBitmap;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
/**
* some check need to do after analyze whole plan.
*/
public class CheckAfterRewrite extends OneAnalysisRuleFactory {
@Override
public Rule build() {
return any().then(plan -> {
checkAllSlotReferenceFromChildren(plan);
checkUnexpectedExpression(plan);
checkMetricTypeIsUsedCorrectly(plan);
checkMatchIsUsedCorrectly(plan);
return null;
}).toRule(RuleType.CHECK_ANALYSIS);
}
private void checkUnexpectedExpression(Plan plan) {
boolean isGenerate = plan instanceof Generate;
boolean isAgg = plan instanceof LogicalAggregate;
boolean isRepeat = plan instanceof LogicalRepeat;
boolean isWindow = plan instanceof LogicalWindow;
boolean notAggAndWindow = !isAgg && !isWindow;
for (Expression expression : plan.getExpressions()) {
expression.foreach(expr -> {
if (expr instanceof SubqueryExpr) {
throw new AnalysisException("Subquery is not allowed in " + plan.getType());
} else if (!isGenerate && expr instanceof TableGeneratingFunction) {
throw new AnalysisException("table generating function is not allowed in " + plan.getType());
} else if (notAggAndWindow && expr instanceof AggregateFunction) {
throw new AnalysisException("aggregate function is not allowed in " + plan.getType());
} else if (!isRepeat && expr instanceof GroupingScalarFunction) {
throw new AnalysisException("grouping scalar function is not allowed in " + plan.getType());
} else if (!isWindow && (expr instanceof WindowExpression || expr instanceof WindowFunction)) {
throw new AnalysisException("analytic function is not allowed in " + plan.getType());
}
});
}
}
/**
*
* used for test
* after applying a rule, use this method to check whether all slots are from children
*/
public void checkTreeAllSlotReferenceFromChildren(Plan plan) {
checkAllSlotReferenceFromChildren(plan);
for (Plan child : plan.children()) {
checkTreeAllSlotReferenceFromChildren(child);
}
}
private void checkAllSlotReferenceFromChildren(Plan plan) {
Set<Slot> inputSlots = plan.getInputSlots();
RoaringBitmap childrenOutput = plan.getChildrenOutputExprIdBitSet();
Set<Slot> notFromChildren = Sets.newHashSet();
for (Slot inputSlot : inputSlots) {
if (!childrenOutput.contains(inputSlot.getExprId().asInt())) {
notFromChildren.add(inputSlot);
}
}
if (notFromChildren.isEmpty()) {
return;
}
if (plan instanceof LogicalRepeat) {
LogicalRepeat repeat = (LogicalRepeat) plan;
notFromChildren.remove(repeat.getGroupingId().get());
}
notFromChildren = removeValidSlotsNotFromChildren(notFromChildren, childrenOutput);
if (!notFromChildren.isEmpty()) {
if (plan.arity() != 0 && plan.child(0) instanceof LogicalAggregate) {
throw new AnalysisException(String.format(plan.getClass().getSimpleName() + " expression %s must"
+ " appear in the GROUP BY clause or be used in an aggregate function",
notFromChildren.stream()
.map(slot -> "'" + slot.getName() + "'")
.collect(Collectors.joining(", "))));
} else {
throw new AnalysisException(String.format(
"Input slot(s) not in child's output: %s in plan: %s\nchild output is: %s\nplan tree:\n%s",
StringUtils.join(notFromChildren.stream().map(ExpressionTrait::toString)
.collect(Collectors.toSet()), ", "),
plan,
plan.children().stream()
.flatMap(child -> child.getOutput().stream())
.collect(Collectors.toSet()),
plan.treeString()));
}
}
}
private Set<Slot> removeValidSlotsNotFromChildren(Set<Slot> slots, RoaringBitmap childrenOutput) {
return slots.stream()
.filter(expr -> {
return !(expr instanceof SlotNotFromChildren);
})
.collect(Collectors.toSet());
}
private void checkMetricTypeIsUsedCorrectly(Plan plan) {
if (!(plan instanceof LogicalJoin)) {
for (Expression expression : plan.getExpressions()) {
checkVariantComparisonIsUsedCorrectly(expression);
}
}
checkVariantV2AggregateFunctions(plan);
if (plan instanceof LogicalAggregate) {
LogicalAggregate<?> agg = (LogicalAggregate<?>) plan;
for (Expression groupBy : agg.getGroupByExpressions()) {
if (groupBy.getDataType().isObjectType()
|| groupBy.getDataType().isVarBinaryType()
|| (groupBy.getDataType().isVariantType() && !isVariantV2Enabled())) {
throw new AnalysisException(Type.OnlyMetricTypeErrorMsg);
}
if (groupBy.getDataType().isVariantType()) {
checkVariantV2ExpressionIsComputed(groupBy);
}
}
} else if (plan instanceof LogicalSort) {
LogicalSort<?> sort = (LogicalSort<?>) plan;
checkTypesSupportOrdering(sort.getOrderKeys().stream()
.map(orderKey -> orderKey.getExpr().getDataType()));
} else if (plan instanceof LogicalTopN) {
LogicalTopN<?> topN = (LogicalTopN<?>) plan;
checkTypesSupportOrdering(topN.getOrderKeys().stream()
.map(orderKey -> orderKey.getExpr().getDataType()));
} else if (plan instanceof LogicalWindow) {
((LogicalWindow<?>) plan).getWindowExpressions().forEach(a -> {
if (!(a instanceof Alias && ((Alias) a).child() instanceof WindowExpression)) {
return;
}
WindowExpression windowExpression = (WindowExpression) ((Alias) a).child();
checkTypesSupportOrdering(Stream.concat(
windowExpression.getOrderKeys().stream().map(orderKey -> orderKey.getDataType()),
windowExpression.getPartitionKeys().stream().map(Expression::getDataType)));
});
} else if (plan instanceof LogicalJoin) {
LogicalJoin<?, ?> join = (LogicalJoin<?, ?>) plan;
for (Expression conjunct : join.getHashJoinConjuncts()) {
if (!isVariantV2Enabled() && conjunct.children().stream()
.anyMatch(child -> child.getDataType().isVariantType())) {
throw new AnalysisException(
"variant type could not in join equal conditions: " + conjunct.toSql());
}
if (conjunct.children().stream()
.anyMatch(child -> child.getDataType().isVariantType())) {
checkVariantV2ExpressionIsComputed(conjunct);
}
if (conjunct.anyMatch(e -> ((Expression) e).getDataType().isVarBinaryType())) {
throw new AnalysisException(
"varbinary type could not in join equal conditions: " + conjunct.toSql());
}
}
for (Expression conjunct : join.getOtherJoinConjuncts()) {
checkVariantComparisonIsUsedCorrectly(conjunct);
if (conjunct.anyMatch(e -> ((Expression) e).getDataType().isVarBinaryType())) {
throw new AnalysisException(
"varbinary type could not in join equal conditions: " + conjunct.toSql());
}
}
for (Expression conjunct : join.getMarkJoinConjuncts()) {
checkVariantComparisonIsUsedCorrectly(conjunct);
if (conjunct.anyMatch(e -> ((Expression) e).getDataType().isVarBinaryType())) {
throw new AnalysisException(
"varbinary type could not in join equal conditions: " + conjunct.toSql());
}
}
} else if (isVariantV2Enabled() && plan instanceof LogicalSetOperation
&& (!(plan instanceof LogicalUnion)
|| ((LogicalSetOperation) plan).getQualifier() == SetOperation.Qualifier.DISTINCT)) {
for (Expression expression : plan.getExpressions()) {
if (expression.getDataType().isVariantType()) {
checkVariantV2ExpressionIsComputed(expression);
}
}
}
}
private boolean isVariantV2Enabled() {
ConnectContext context = ConnectContext.get();
return context != null && context.getSessionVariable().isEnableVariantV2();
}
private void checkVariantV2AggregateFunctions(Plan plan) {
if (!isVariantV2Enabled()) {
return;
}
plan.getExpressions().forEach(expression -> expression.foreach(expr -> {
if (!(expr instanceof AggregateFunction)
|| (expr instanceof Count && !((Count) expr).isDistinct())) {
return;
}
AggregateFunction function = (AggregateFunction) expr;
if (function.getArguments().stream()
.anyMatch(argument -> argument.getDataType().containsVariantType())) {
throw new AnalysisException(function.getName()
+ " does not support Variant arguments when enable_variant_v2 is true");
}
}));
}
private void checkVariantV2ExpressionIsComputed(Expression expression) {
if (expression.anyMatch(expr -> expr instanceof SlotReference
&& ((SlotReference) expr).getDataType().isVariantType()
&& ((SlotReference) expr).getOriginalColumn().isPresent())) {
throw new AnalysisException("Variant V2 equality currently supports computed Variant "
+ "expressions only; cast stored Variant values to a concrete type first");
}
}
private void checkTypesSupportOrdering(Stream<DataType> dataTypes) {
List<DataType> types = dataTypes.collect(Collectors.toList());
if (types.stream().anyMatch(DataType::containsVariantType)) {
throw new AnalysisException(VariantType.UNSUPPORTED_ORDERING_COMPARISON_MESSAGE);
}
if (types.stream().anyMatch(DataType::isObjectOrVariantType)) {
throw new AnalysisException(Type.OnlyMetricTypeErrorMsg);
}
}
private void checkVariantComparisonIsUsedCorrectly(Expression expression) {
expression.foreach(expr -> {
if (expr instanceof ComparisonPredicate) {
ComparisonPredicate comparisonPredicate = (ComparisonPredicate) expr;
boolean leftIsVariant = comparisonPredicate.left().getDataType().isVariantType();
boolean rightIsVariant = comparisonPredicate.right().getDataType().isVariantType();
if (leftIsVariant || rightIsVariant) {
DataType variantDataType = leftIsVariant
? comparisonPredicate.left().getDataType()
: comparisonPredicate.right().getDataType();
throw new AnalysisException("data type " + variantDataType
+ " could not used in ComparisonPredicate " + comparisonPredicate.toSql()
+ ". " + VariantType.UNSUPPORTED_ORDERING_COMPARISON_MESSAGE);
}
}
});
}
private void checkMatchIsUsedCorrectly(Plan plan) {
for (Expression expression : plan.getExpressions()) {
if (expression instanceof Match) {
if (plan instanceof LogicalFilter && plan.child(0) instanceof LogicalOlapScan) {
return;
} else {
throw new AnalysisException(String.format(
"Not support match in %s in plan: %s, only support in olapScan filter",
plan.child(0), plan));
}
}
}
}
}