airborne12 commented on code in PR #69292: URL: https://github.com/apache/doris/pull/69292#discussion_r4345863334
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindScore.java: ########## @@ -0,0 +1,440 @@ +// 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.Column; +import org.apache.doris.catalog.FunctionRegistry; +import org.apache.doris.nereids.analyzer.UnboundAlias; +import org.apache.doris.nereids.analyzer.UnboundFunction; +import org.apache.doris.nereids.analyzer.UnboundRelation; +import org.apache.doris.nereids.analyzer.UnboundSlot; +import org.apache.doris.nereids.analyzer.UnboundStar; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.properties.OrderKey; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.ComparisonPredicate; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.GreaterThan; +import org.apache.doris.nereids.trees.expressions.GreaterThanEqual; +import org.apache.doris.nereids.trees.expressions.LessThan; +import org.apache.doris.nereids.trees.expressions.LessThanEqual; +import org.apache.doris.nereids.trees.expressions.Match; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SubqueryExpr; +import org.apache.doris.nereids.trees.expressions.WindowExpression; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +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.LogicalCTE; +import org.apache.doris.nereids.trees.plans.logical.LogicalCheckPolicy; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalHaving; +import org.apache.doris.nereids.trees.plans.logical.LogicalLimit; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalPreAggOnHint; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalQualify; +import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat; +import org.apache.doris.nereids.trees.plans.logical.LogicalSelectHint; +import org.apache.doris.nereids.trees.plans.logical.LogicalSetOperation; +import org.apache.doris.nereids.trees.plans.logical.LogicalSink; +import org.apache.doris.nereids.trees.plans.logical.LogicalSort; +import org.apache.doris.nereids.trees.plans.logical.LogicalSubQueryAlias; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * Bind score() to the olap scan that computes it. + * + * <p>score() is the relevance of a row for the MATCH and SEARCH predicates in the WHERE clause of + * its SELECT, and only the olap scan that evaluates those predicates computes it. So score() is + * supported in a SELECT that reads one olap table and has a MATCH or SEARCH predicate in its WHERE + * clause: + * <ul> + * <li>in the WHERE clause, as one minimum score: score() > x or score() >= x;</li> + * <li>in the select list, window functions and ORDER BY of a SELECT that does not aggregate.</li> + * </ul> + * A SELECT in a subquery, view, CTE or set operation has to ORDER BY score() with a LIMIT, which keeps + * the predicates of the outer query away from its scan. A correlated subquery cannot call score(), see + * CheckScoreInCorrelatedSubquery. Any other use of score() is rejected. + * + * <p>This rule runs first on each plan the analyzer binds: a statement, and on their own its views, CTEs + * and subqueries. The plan nodes of each SELECT are still its clauses then. The rule asks the table of a + * SELECT that calls score() for a score column, which BindRelation adds to the olap scan, and replaces + * the score() calls of the SELECT with that column. The scan scores every row it returns, so the minimum + * score is a filter on the column. PushDownScoreTopNIntoOlapScan pushes a TopN on the score into the + * scan. + */ +public class BindScore implements CustomRewriter { + public static final String UNSUPPORTED_SELECT_ERROR = "score() can only be used in a SELECT that reads one " + + "table and has a MATCH or SEARCH predicate in its WHERE clause"; + public static final String NESTED_SELECT_ERROR = "score() in a subquery, view, CTE or set operation " + + "requires ORDER BY score() with LIMIT"; + public static final String CORRELATED_SUBQUERY_ERROR = "score() cannot be used in a correlated subquery"; + public static final String AGGREGATE_ERROR = "score() can only be used in the WHERE clause of a SELECT " + + "with GROUP BY, aggregate functions or DISTINCT"; + public static final String HAVING_OR_QUALIFY_ERROR = "score(), MATCH and SEARCH cannot be used in HAVING " + + "or QUALIFY of a SELECT that calls score()"; + public static final String ONE_RANGE_ERROR = "Only one score() range predicate is supported in WHERE clause"; + public static final String RANGE_FORM_ERROR = "score() predicate in WHERE clause must be in the form of " + + "'score() > literal' or 'score() >= literal'"; + public static final String TOP_LEVEL_RANGE_ERROR = "score() predicate must be a top-level AND condition " + + "in WHERE clause"; + public static final String NUMERIC_COMPARISON_ERROR = "score() function can only be compared with numeric types"; + + // No column of a table can have a name with the hidden column prefix. + private static final String SCORE_COLUMN = Column.HIDDEN_COLUMN_PREFIX + "SCORE__"; + + @Override + public Plan rewriteRoot(Plan plan, JobContext jobContext) { + if (!plan.anyMatch(node -> callsScore((Plan) node))) { + return plan; + } + FunctionRegistry functions = jobContext.getCascadesContext().getConnectContext().getEnv() + .getFunctionRegistry(); + // A view, a CTE and a subquery are analyzed without the sink of their statement. + Plan newPlan = bindQuery(plan, plan instanceof LogicalSink, functions); + // Outside the clauses of a SELECT, e.g. in a join condition. + if (newPlan.anyMatch(node -> callsScore((Plan) node))) { + throw new AnalysisException(UNSUPPORTED_SELECT_ERROR); + } + return newPlan; + } + + /** The score column that BindRelation added to the scan. */ + public static Optional<Slot> findScoreSlot(LogicalOlapScan scan) { + return scan.getVirtualColumns().stream() + .filter(column -> column instanceof Alias && ((Alias) column).child() instanceof Score) + .map(NamedExpression::toSlot) + .findFirst(); + } + + /** Make the scan of a table that a SELECT asked for its score column score every row it returns. */ + public static LogicalOlapScan addScoreColumn(LogicalOlapScan scan, Alias scoreColumn) { + // A zero score limit scores every row. + return scan.appendVirtualColumnsAndTopN(ImmutableList.of(scoreColumn), scan.getAnnOrderKeys(), + scan.getAnnLimit(), ImmutableList.of(new OrderKey(scoreColumn.toSlot(), false, false)), + Optional.of(0L), Optional.empty()); + } + + /** Only the olap scan of a table can compute the score, not a view, a CTE or another relation. */ + public static void checkScoreColumn(Plan relation, Alias scoreColumn) { + if (!relation.getOutputSet().contains(scoreColumn.toSlot())) { + throw new AnalysisException(UNSUPPORTED_SELECT_ERROR); + } + } + + // A query: a SELECT, a set operation of queries or a query with a WITH clause, also under the sink of + // the statement. + private static Plan bindQuery(Plan plan, boolean outermost, FunctionRegistry functions) { + if (plan instanceof LogicalSink || plan instanceof LogicalCTE) { + return withChildren(plan, ImmutableList.of(bindQuery(plan.child(0), outermost, functions))); + } + if (plan instanceof LogicalSetOperation) { + return withChildren(plan, plan.children().stream() + .map(child -> bindQuery(child, false, functions)) + .collect(ImmutableList.toImmutableList())); + } + return bindSelect(plan, outermost, functions); + } + + private static Plan bindSelect(Plan top, boolean outermost, FunctionRegistry functions) { + // The parser builds a SELECT as the chain of its clauses, from its hints down to its WHERE clause. + List<Plan> clauses = new ArrayList<>(); + Plan from = top; + for (; isClause(from); from = from.child(0)) { + clauses.add(from); + } + Plan newFrom = bindFrom(from, functions); + if (clauses.stream().anyMatch(BindScore::callsScore)) { + check(clauses, from, outermost, functions); + Alias scoreColumn = new Alias(new Score(), SCORE_COLUMN); + newFrom = askForScoreColumn(from, scoreColumn); + clauses = clauses.stream().map(clause -> bindScore(clause, scoreColumn.toSlot())) + .collect(Collectors.toList()); + } + Plan plan = newFrom; + for (int i = clauses.size() - 1; i >= 0; i--) { + plan = withChildren(clauses.get(i), ImmutableList.of(plan)); + } + return plan; + } + + private static boolean isClause(Plan plan) { + return plan instanceof LogicalPreAggOnHint || plan instanceof LogicalSelectHint + || plan instanceof LogicalLimit || plan instanceof LogicalSort || plan instanceof LogicalQualify + || plan instanceof LogicalHaving || plan instanceof LogicalProject + || plan instanceof LogicalAggregate || plan instanceof LogicalRepeat || plan instanceof LogicalFilter; + } + + // The FROM clause: tables, joins and lateral views of them, and subqueries, which are queries of + // their own. + private static Plan bindFrom(Plan from, FunctionRegistry functions) { + if (from.arity() == 0 || isTable(from)) { + return from; + } + if (from instanceof LogicalSubQueryAlias) { + return withChildren(from, ImmutableList.of(bindQuery(from.child(0), false, functions))); + } + // E.g. a set operation under ORDER BY and LIMIT. + if (isClause(from) || from instanceof LogicalSetOperation || from instanceof LogicalCTE) { + return bindQuery(from, false, functions); + } + return withChildren(from, from.children().stream() + .map(child -> bindFrom(child, functions)) + .collect(ImmutableList.toImmutableList())); + } + + // A table, also with an alias. + private static boolean isTable(Plan plan) { + Plan relation = plan instanceof LogicalSubQueryAlias ? plan.child(0) : plan; + return relation instanceof LogicalCheckPolicy && relation.child(0) instanceof UnboundRelation; + } + + private static Plan askForScoreColumn(Plan table, Alias scoreColumn) { + if (table instanceof LogicalSubQueryAlias) { + return table.withChildren(askForScoreColumn(table.child(0), scoreColumn)); + } + return table.withChildren(((UnboundRelation) table.child(0)).withScoreColumn(scoreColumn)); + } + + private static void check(List<Plan> clauses, Plan from, boolean outermost, FunctionRegistry functions) { + // A subquery in a clause reads another table. + if (!isTable(from) || clauses.stream().anyMatch(clause -> clause.getExpressions().stream() + .anyMatch(expression -> expression.containsType(SubqueryExpr.class)))) { + throw new AnalysisException(UNSUPPORTED_SELECT_ERROR); + } + Plan last = clauses.get(clauses.size() - 1); + List<Expression> where = last instanceof LogicalFilter + ? ((LogicalFilter<?>) last).getConjuncts().stream() + .flatMap(conjunct -> ExpressionUtils.extractConjunction(conjunct).stream()) + .collect(Collectors.toList()) + : ImmutableList.of(); + if (where.stream().noneMatch(conjunct -> conjunct.anyMatch(BindScore::matchesText))) { + throw new AnalysisException(UNSUPPORTED_SELECT_ERROR); + } + clauses.forEach(clause -> clause.getExpressions().forEach(expression -> expression.foreach(node -> { + if (node instanceof ComparisonPredicate) { + checkComparedWithNumber((ComparisonPredicate) node); + } + }))); + List<Expression> ranges = where.stream() + .filter(conjunct -> conjunct.anyMatch(BindScore::isScoreCall)) + .collect(Collectors.toList()); + if (ranges.size() > 1) { + throw new AnalysisException(ONE_RANGE_ERROR); + } + ranges.forEach(BindScore::checkRange); + + List<Plan> aboveWhere = last instanceof LogicalFilter ? clauses.subList(0, clauses.size() - 1) : clauses; + boolean aggregates = aboveWhere.stream().anyMatch(clause -> aggregates(clause, functions)); + // HAVING and QUALIFY can also read score(), MATCH and SEARCH through aliases of the select list. + Set<String> scoringAliases = new HashSet<>(); + for (int i = aboveWhere.size() - 1; i >= 0; i--) { + for (Expression expression : aboveWhere.get(i).getExpressions()) { + if (expression instanceof UnboundAlias && ((UnboundAlias) expression).getAlias().isPresent() + && scores(expression, scoringAliases)) { + scoringAliases.add(((UnboundAlias) expression).getAlias().get().toLowerCase()); + } + } + } + for (Plan clause : aboveWhere) { + if ((clause instanceof LogicalHaving || clause instanceof LogicalQualify) + && clause.getExpressions().stream().anyMatch(expression -> scores(expression, scoringAliases))) { + throw new AnalysisException(HAVING_OR_QUALIFY_ERROR); + } + if (aggregates && callsScore(clause)) { + throw new AnalysisException(AGGREGATE_ERROR); + } + } + if (!outermost && !ordersByScoreWithLimit(aboveWhere)) { + throw new AnalysisException(NESTED_SELECT_ERROR); + } + } + + // score() is a FLOAT. The type of a literal is known before binding, e.g. score() > 'abc' is rejected. + private static void checkComparedWithNumber(ComparisonPredicate comparison) { + for (int i = 0; i < 2; i++) { + Expression other = comparison.child(1 - i); + if (isScoreCall(comparison.child(i)) && other instanceof Literal + && !other.getDataType().isNumericType() && !other.getDataType().isNullType()) { + throw new AnalysisException(NUMERIC_COMPARISON_ERROR + ", but found: " + other.getDataType() + + " in " + comparison.toSql()); + } + } + } + + // score() > x or score() >= x, also written x < score() or x <= score(). A cast of score() to + // FLOAT or DOUBLE keeps its value. + private static void checkRange(Expression conjunct) { + if (!(conjunct instanceof ComparisonPredicate)) { + throw new AnalysisException(TOP_LEVEL_RANGE_ERROR); + } + ComparisonPredicate comparison = (ComparisonPredicate) conjunct; + boolean scoreAbove = (comparison instanceof GreaterThan || comparison instanceof GreaterThanEqual) + && isScore(comparison.left()) && isConstant(comparison.right()); + boolean scoreBelow = (comparison instanceof LessThan || comparison instanceof LessThanEqual) + && isScore(comparison.right()) && isConstant(comparison.left()); + if (!scoreAbove && !scoreBelow) { + throw new AnalysisException(RANGE_FORM_ERROR); + } + } + + private static boolean isScore(Expression expression) { + if (expression instanceof Cast) { + return expression.getDataType().isFloatLikeType() && isScore(((Cast) expression).child()); + } + return isScoreCall(expression); + } + + private static boolean isConstant(Expression expression) { + return expression.getInputSlots().isEmpty() && !expression.anyMatch(BindScore::isScoreCall); + } + + // GROUP BY, DISTINCT or an aggregate function, which aggregates all rows of a SELECT without GROUP BY. + private static boolean aggregates(Plan clause, FunctionRegistry functions) { + return clause instanceof LogicalAggregate || clause instanceof LogicalRepeat + || clause instanceof LogicalProject && ((LogicalProject<?>) clause).isDistinct() + || clause.getExpressions().stream().anyMatch(expression -> callsAggregate(expression, functions)); + } + + private static boolean callsAggregate(Expression expression, FunctionRegistry functions) { + // A window function aggregates the rows of its window. + if (expression instanceof WindowExpression) { + return false; + } + if (expression instanceof UnboundFunction && functions.isAggregateFunction( + ((UnboundFunction) expression).getDbName(), ((UnboundFunction) expression).getName())) { + return true; + } + return expression.children().stream().anyMatch(child -> callsAggregate(child, functions)); + } + + private static boolean scores(Expression expression, Set<String> scoringAliases) { + return expression.anyMatch(node -> isScoreCall((Expression) node) || matchesText((Expression) node) + || node instanceof UnboundSlot && ((UnboundSlot) node).getNameParts().size() == 1 + && scoringAliases.contains(((UnboundSlot) node).getName().toLowerCase())); + } + + // A TopN on the score at the top of the SELECT, which no outer predicate passes. + private static boolean ordersByScoreWithLimit(List<Plan> aboveWhere) { + for (int i = 0; i + 1 < aboveWhere.size(); i++) { + if (aboveWhere.get(i) instanceof LogicalLimit && aboveWhere.get(i + 1) instanceof LogicalSort) { + List<OrderKey> orderKeys = ((LogicalSort<?>) aboveWhere.get(i + 1)).getOrderKeys(); + return orderKeys.size() == 1 && isScoreOrAlias(orderKeys.get(0).getExpr(), aboveWhere); + } + } + return false; + } + + private static boolean isScoreOrAlias(Expression expression, List<Plan> aboveWhere) { + if (isScoreCall(expression)) { + return true; + } + return expression instanceof UnboundSlot && ((UnboundSlot) expression).getNameParts().size() == 1 + && aboveWhere.stream().flatMap(clause -> clause.getExpressions().stream()) + .anyMatch(item -> item instanceof UnboundAlias && isScoreCall(item.child(0)) + && ((UnboundAlias) item).getAlias() + .filter(((UnboundSlot) expression).getName()::equalsIgnoreCase).isPresent()); + } + + // Replaces the score() calls of a clause with the score column, which * leaves out because it is not + // a column of the table. + private static Plan bindScore(Plan clause, Slot scoreSlot) { + Function<Expression, Expression> bind = expression -> expression.rewriteDownShortCircuit( + node -> isScoreCall(node) ? scoreSlot : node); + if (clause instanceof LogicalFilter) { + LogicalFilter<?> filter = (LogicalFilter<?>) clause; + return filter.withConjuncts(filter.getConjuncts().stream().map(bind) + .collect(ImmutableSet.toImmutableSet())); + } else if (clause instanceof LogicalSort) { + LogicalSort<?> sort = (LogicalSort<?>) clause; + return sort.withOrderKeys(sort.getOrderKeys().stream() + .map(orderKey -> orderKey.withExpression(bind.apply(orderKey.getExpr()))) + .collect(ImmutableList.toImmutableList())); + } else if (clause instanceof LogicalProject) { + LogicalProject<?> project = (LogicalProject<?>) clause; + return project.withProjects(project.getProjects().stream() + .map(item -> exceptScore((NamedExpression) bind.apply(item), scoreSlot)) + .collect(ImmutableList.toImmutableList())); + } else if (clause instanceof LogicalAggregate) { + LogicalAggregate<?> aggregate = (LogicalAggregate<?>) clause; + return aggregate.withAggOutput(aggregate.getOutputExpressions().stream() + .map(item -> exceptScore(item, scoreSlot)) + .collect(ImmutableList.toImmutableList())); + } else if (clause instanceof LogicalRepeat) { + LogicalRepeat<?> repeat = (LogicalRepeat<?>) clause; + return repeat.withAggOutput(repeat.getOutputExpressions().stream() + .map(item -> exceptScore(item, scoreSlot)) + .collect(ImmutableList.toImmutableList())); + } + return clause; + } + + private static NamedExpression exceptScore(NamedExpression item, Slot scoreSlot) { + if (!(item instanceof UnboundStar)) { + return item; + } + UnboundStar star = (UnboundStar) item; + return new UnboundStar(star.getQualifier(), ImmutableList.<NamedExpression>builder() + .addAll(star.getExceptedSlots()).add(scoreSlot).build(), + star.getReplacedAlias(), star.getIndexInSqlString()); + } + + private static boolean callsScore(Plan plan) { + return plan.getExpressions().stream().anyMatch(expression -> expression.anyMatch(BindScore::isScoreCall)); Review Comment: Update: 854a54f23e moves the binding back to the start of rewrite (the analysis classes are as on master again). The binder now runs after analysis expanded `* REPLACE` and alias functions, so `score()` in both is bound like any select item instead of being rejected; `BindScoreToOlapScanTest.testStarLeavesOutTheScoreColumn` and `testScoreInAliasFunction` cover them. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
