morrySnow commented on code in PR #66288: URL: https://github.com/apache/doris/pull/66288#discussion_r3773699426
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/InferPredicateFromMonotonicFunction.java: ########## @@ -0,0 +1,356 @@ +// 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.expression.rules; + +import org.apache.doris.nereids.trees.expressions.And; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.ComparisonPredicate; +import org.apache.doris.nereids.trees.expressions.EqualTo; +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.Or; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.functions.DateCeilFloorMonotonic; +import org.apache.doris.nereids.trees.expressions.functions.Monotonic; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Left; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Substring; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Year; +import org.apache.doris.nereids.trees.expressions.literal.DateV2Literal; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral; +import org.apache.doris.nereids.types.coercion.CharacterType; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.TypeCoercionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.common.collect.Sets; + +import java.util.List; +import java.util.Optional; +import java.util.Set; + +// When a partition predicate wraps the partition column in a function +// (e.g. date_trunc(dt, 'day') >= '2020-01-01'), this class derives an extra predicate on the +// bare column (dt >= '2020-01-01') and appends it, so partition pruning has a range to work with. +// +// Why: the fast pruning path binary-searches a partition list sorted by the partition column +// (PartitionPredicateToRange). Binary search only understands "bare column op literal" shapes that +// map to a range; once the column is wrapped by date_trunc / year / substring it cannot, and pruning +// falls back to per-partition evaluation. The derived predicate is the key that unlocks binary +// search. It lives only inside pruning and is never written back as a runtime filter, so it adds no +// per-row cost. See the caller PartitionPruner.pruneInternal. +// +// input: date_trunc(dt, 'day') >= '2020-01-05' (date_trunc is a floor, floor(x) <= x) +// infer: floor(dt) >= '2020-01-05' => dt >= '2020-01-05' +// output: date_trunc(dt, 'day') >= '2020-01-05' AND dt >= '2020-01-05' +// +// Every derived predicate is a necessary condition of the original, resting only on the function's +// monotonic/rounding property and never on data distribution, so appending it never changes results. +// Three kinds are supported: +// - prefix: substring(col,1,n) >= 'abc' => col >= 'abc' (a prefix never sorts after the whole string) +// - rounding: floor(col) >= c => col >= c; ceil(col) <= c => col <= c +// - year: year(col) op y => a boundary date or the half-open range [y-01-01, (y+1)-01-01) +// For each kind, EqualTo yields only a single-sided non-strict bound (floor(col)=c gives col >= c only). +// +// from_unixtime and similar are excluded on purpose: they map a unix timestamp to local time, so the +// inverse depends on timezone shifts (DST) rather than a globally fixed property, and no safe bare-column +// bound can be derived. +final class InferPredicateFromMonotonicFunction { + private InferPredicateFromMonotonicFunction() { + } + + // Comparison: derive directly. AND/OR: recurse into children and keep the boolean structure. + // NOT and other shapes are returned unchanged, since a necessary condition cannot cross negation. + static Expression inferForPartitionPrune(Expression partitionPredicate) { + if (partitionPredicate instanceof ComparisonPredicate) { + return appendInferredPredicate((ComparisonPredicate) partitionPredicate); + } + if (partitionPredicate instanceof And) { + return rewriteConjunction((And) partitionPredicate); + } + if (partitionPredicate instanceof Or) { + return rewriteDisjunction((Or) partitionPredicate); + } + return partitionPredicate; + } + + // Append the derived bare-column predicate with AND; return the comparison unchanged if nothing is derived. + private static Expression appendInferredPredicate(ComparisonPredicate comparison) { + Optional<Expression> inferred = infer(comparison); + if (!inferred.isPresent()) { + return comparison; + } + // year EqualTo derives a lower AND upper bound; flatten so everything joins one AND. + List<Expression> conjuncts = Lists.newArrayList(comparison); + conjuncts.addAll(ExpressionUtils.extractConjunction(inferred.get())); + return ExpressionUtils.and(conjuncts); + } + + // Recurse into each conjunct and flatten derived bare-column predicates into the current AND, deduping. + private static Expression rewriteConjunction(And predicate) { + List<Expression> conjuncts = ExpressionUtils.extractConjunction(predicate); + Set<Expression> existing = Sets.newHashSet(); + List<Expression> rewrittenConjuncts = Lists.newArrayList(); + for (Expression conjunct : conjuncts) { + Expression rewritten = inferForPartitionPrune(conjunct); + for (Expression rewrittenConjunct : ExpressionUtils.extractConjunction(rewritten)) { + if (existing.add(rewrittenConjunct)) { + rewrittenConjuncts.add(rewrittenConjunct); + } + } + } + return predicate.withChildren(rewrittenConjuncts); + } + + // Derive per branch and keep the OR: necessary conditions of different branches must not be + // lifted out and ANDed together. PartitionPredicateToRange unions the branch ranges. + private static Expression rewriteDisjunction(Or predicate) { + List<Expression> children = predicate.children(); + List<Expression> rewrittenChildren = Lists.newArrayListWithCapacity(children.size()); + for (Expression child : children) { + Expression rewritten = inferForPartitionPrune(child); + rewrittenChildren.add(rewritten); + } + return predicate.withChildren(rewrittenChildren); + } + + // Try the three inference kinds in order. Precondition: the right side must be a literal. + private static Optional<Expression> infer(ComparisonPredicate comparison) { + if (!(comparison.right() instanceof Literal)) { + return Optional.empty(); + } + + return inferPrefixPredicate(comparison) + .or(() -> inferYearPredicate(comparison)) + .or(() -> inferRoundingPredicate(comparison)); + } + + // Prefix inference: a prefix never sorts after the whole string (prefix(s) <= s), so + // prefix(col) >= 'abc' => col >= 'abc'. >/>= keep the operator; = yields only the lower bound; + // </<= do not hold (prefix(col) <= 'abc' allows col = 'abd', which is larger). + private static Optional<Expression> inferPrefixPredicate(ComparisonPredicate comparison) { + if (!(comparison.right() instanceof StringLikeLiteral)) { + return Optional.empty(); + } + Optional<Expression> source = prefixSource(comparison.left()); + if (!source.isPresent()) { + return Optional.empty(); + } + + ComparisonPredicate inferred; + if (comparison instanceof GreaterThan || comparison instanceof GreaterThanEqual) { + inferred = (ComparisonPredicate) comparison.withChildren(source.get(), comparison.right()); + } else if (comparison instanceof EqualTo) { + inferred = new GreaterThanEqual(source.get(), comparison.right()); + } else { + return Optional.empty(); + } + return Optional.of(inferredPredicate(inferred)); + } + + // Extract the bare column from a fixed-length, start-anchored prefix over a character Slot: + // substring(col, 1, n) (position 1, length present) or left(col, n). Otherwise empty. + private static Optional<Expression> prefixSource(Expression expression) { + Expression source; + Expression length; + if (expression instanceof Substring) { + Substring substring = (Substring) expression; + if (!(substring.getPosition() instanceof IntegerLikeLiteral) + || ((IntegerLikeLiteral) substring.getPosition()).getIntValue() != 1 + || !substring.getLength().isPresent()) { + return Optional.empty(); + } + source = substring.getSource(); + length = substring.getLength().get(); + } else if (expression instanceof Left) { + source = expression.child(0); + length = expression.child(1); + } else { + return Optional.empty(); + } + + // A clean fixed-length prefix requires a character Slot source and a positive integer length. + if (!(source instanceof Slot) || !(source.getDataType() instanceof CharacterType) + || !(length instanceof IntegerLikeLiteral) + || ((IntegerLikeLiteral) length).getIntValue() <= 0) { + return Optional.empty(); + } + return Optional.of(source); + } + + // Rounding inference over a bare date column. floor(x) <= x gives a lower bound; + // ceil(x) >= x gives an upper bound. EqualTo yields only the single-sided non-strict bound. + // floor(dt) >= '2020-01-05' => dt >= '2020-01-05' + // ceil(dt) <= '2020-01-05' => dt <= '2020-01-05' + private static Optional<Expression> inferRoundingPredicate(ComparisonPredicate comparison) { + if (!(comparison.left() instanceof Monotonic)) { + return Optional.empty(); + } + Monotonic function = (Monotonic) comparison.left(); + // A rounding function may take several arguments (e.g. date_trunc(dt, 'day')); + // getMonotonicFunctionChildIndex points at the rounded input column. + Expression source = function.child(function.getMonotonicFunctionChildIndex()); + if (!isDateSlot(source) || !hasSafeRoundingArguments(function)) { + return Optional.empty(); + } + + ComparisonPredicate inferred; + if (function.isFloor()) { + if (comparison instanceof GreaterThan || comparison instanceof GreaterThanEqual) { + inferred = (ComparisonPredicate) comparison.withChildren(source, comparison.right()); + } else if (comparison instanceof EqualTo) { + inferred = new GreaterThanEqual(source, comparison.right()); + } else { + // floor gives no upper bound: floor(dt) <= c allows dt beyond c within the same bucket. + return Optional.empty(); + } + } else if (function.isCeil()) { + if (comparison instanceof LessThan || comparison instanceof LessThanEqual) { + inferred = (ComparisonPredicate) comparison.withChildren(source, comparison.right()); + } else if (comparison instanceof EqualTo) { + inferred = new LessThanEqual(source, comparison.right()); + } else { + return Optional.empty(); + } + } else { + return Optional.empty(); + } + return Optional.of(inferredPredicate(inferred)); + } + + // Verify the period/origin arguments do not break floor(x) <= x / ceil(x) >= x. + // Only DateCeilFloorMonotonic (day_ceil/floor, hour_ceil/floor, ...) carries such arguments; + // everything else (date_trunc, to_date, ...) is safe by construction. + private static boolean hasSafeRoundingArguments(Monotonic function) { + if (!(function instanceof DateCeilFloorMonotonic)) { + return true; + } + DateCeilFloorMonotonic roundingFunction = (DateCeilFloorMonotonic) function; + if (roundingFunction.isFloor()) { + return hasValidFloorPeriod(roundingFunction); + } + if (roundingFunction.isCeil()) { + return hasSafeCeilArguments(roundingFunction); + } + return false; + } + + // floor(x) <= x holds for any origin, so a custom origin (the 3-arg form) is allowed; only a + // non-positive explicit period would break it. The 2-arg-with-non-integer case is date_floor(dt, + // interval 3 day) before binding: the interval, not a period, sits at child(1), so it is safe. + private static boolean hasValidFloorPeriod(DateCeilFloorMonotonic function) { + if (function.arity() == 1) { + return true; + } + if (function.arity() == 2 && !function.child(1).getDataType().isIntegerLikeType()) { + return true; + } + return (function.arity() == 2 || function.arity() == 3) + && isPositiveIntegerLiteral(function.child(1)); + } + + // ceil(x) >= x can break under a custom origin, so only origin-free forms are inferred: + // the 1-arg form, or the 2-arg form with a positive integer period. + private static boolean hasSafeCeilArguments(DateCeilFloorMonotonic function) { + return function.arity() == 1 + || (function.arity() == 2 && isPositiveIntegerLiteral(function.child(1))); + } + + private static boolean isPositiveIntegerLiteral(Expression expression) { + return expression instanceof IntegerLikeLiteral + && ((IntegerLikeLiteral) expression).getIntValue() > 0; + } + + // Year inference: map year(dt) op y to a range on dt, where start = y-01-01, end = (y+1)-01-01, + // so year(dt) = y is the half-open range [start, end). Operator mapping (y=2020 example): + // year(dt) >= 2020 -> dt >= 2020-01-01 + // year(dt) > 2020 -> dt >= 2021-01-01 (after year y means not before next year's start) + // year(dt) < 2020 -> dt < 2020-01-01 + // year(dt) <= 2020 -> dt < 2021-01-01 (not after year y means before next year's start) + // year(dt) = 2020 -> dt >= 2020-01-01 AND dt < 2021-01-01 + private static Optional<Expression> inferYearPredicate(ComparisonPredicate comparison) { + // year returns SmallInt, so comparing against a wider integer literal wraps it in an integer + // cast (cast(year(dt) as INT)); such a widening cast is value-preserving, so see through it. + Expression left = comparison.left(); + if (left instanceof Cast && left.child(0) instanceof Year + && left.getDataType().isIntegerLikeType()) { Review Comment: The cast see-through here accepts *any* integer target type, but only a widening cast is value-preserving. `year()` returns SMALLINT, so `cast(year(dt) as INT/BIGINT/LARGEINT)` is safe, but a narrowing cast such as `cast(year(dt) as TINYINT)` truncates/wraps the year. For example `cast(year(dt) as TINYINT) = 5` is also satisfied by rows in years 261, 517, ... (mod 256), yet this branch derives `dt >= 0005-01-01 AND dt < 0006-01-01`, which would incorrectly prune those partitions. Suggest restricting the see-through to cast targets at least as wide as SMALLINT (e.g. compare the target type's width/size against SmallIntType). Low severity: it requires an explicit narrowing cast, which is unusual in practice. -- 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]
