[BugFix] RangePredicateExtractor discard predicates mistakenly (#50854)

Signed-off-by: satanson <[email protected]> (cherry picked from commit cffaa19) # Conflicts: # fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarRangePredicateExtractor.java # fe/fe-core/src/test/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorRewriterTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/PruneUKFKJoinRuleTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/SubqueryTest.java
StarRocks · Sep 10, 2024 · 67498d5 · 67498d5
1 parent 7207d3a
commit 67498d5
Show file tree

Hide file tree

Showing 4 changed files with 756 additions and 1 deletion.
diff --git a/...core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarRangePredicateExtractor.java b/...core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarRangePredicateExtractor.java
@@ -53,6 +53,7 @@ private ScalarOperator rewrite(ScalarOperator predicate, boolean onlyExtractColu
 
         Set<ScalarOperator> conjuncts = Sets.newLinkedHashSet();
         conjuncts.addAll(Utils.extractConjuncts(predicate));
+        predicate = Utils.compoundAnd(conjuncts);
 
         Map<ScalarOperator, ValueDescriptor> extractMap = extractImpl(predicate);
 
@@ -78,7 +79,6 @@ private ScalarOperator rewrite(ScalarOperator predicate, boolean onlyExtractColu
             return predicate;
         }
 
-        predicate = Utils.compoundAnd(Lists.newArrayList(conjuncts));
         if (isOnlyOrCompound(predicate)) {
             Set<ColumnRefOperator> c = Sets.newHashSet(Utils.extractColumnRef(predicate));
             if (c.size() == extractMap.size() &&
@@ -91,7 +91,22 @@ private ScalarOperator rewrite(ScalarOperator predicate, boolean onlyExtractColu
             List<ScalarOperator> cs = Utils.extractConjuncts(predicate);
             Set<ColumnRefOperator> cf = new HashSet<>(Utils.extractColumnRef(predicate));
 
+<<<<<<< HEAD
             if (extractMap.values().stream().allMatch(valueDescriptor -> valueDescriptor.sourceCount == cs.size())
+=======
+            // getSourceCount = cs.size() means that all and components have the same column ref
+            // and it can be merged into one range predicate. mistakenly, when the predicate is
+            // date_trunc(YEAR, dt) = '2024-01-01' AND mode = 'Buzz' AND
+            // date_trunc(YEAR, dt) = '2024-01-01' AND mode = 'Buzz' //duplicate
+            //
+            // only mode = 'Buzz' is a extractable range predicate, so its corresponding ValueDescriptor's
+            // sourceCount = 2(since it occurs twice), and cs(it is also 2 in this example)are number of
+            // unique column refs of the predicate, the two values are properly equivalent, so it yields
+            // wrong result.
+            //
+            // Components of AND/OR should be deduplicated at first to avoid this issue.
+            if (extractMap.values().stream().allMatch(valueDescriptor -> valueDescriptor.getSourceCount() == cs.size())
+>>>>>>> cffaa196d9 ([BugFix] RangePredicateExtractor discard predicates mistakenly (#50854))
                     && extractMap.size() == cf.size()) {
                 return extractExpr;
             }

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorRewriterTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorRewriterTest.java
@@ -2,11 +2,20 @@
 
 package com.starrocks.sql.optimizer.rewrite;
 
+import com.google.common.base.Preconditions;
+import com.google.common.base.Supplier;
 import com.google.common.collect.Lists;
+<<<<<<< HEAD
+=======
+import com.starrocks.analysis.BinaryType;
+import com.starrocks.catalog.FunctionSet;
+>>>>>>> cffaa196d9 ([BugFix] RangePredicateExtractor discard predicates mistakenly (#50854))
 import com.starrocks.catalog.Type;
+import com.starrocks.sql.optimizer.Utils;
 import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.scalar.BetweenPredicateOperator;
 import com.starrocks.sql.optimizer.operator.scalar.BinaryPredicateOperator;
+import com.starrocks.sql.optimizer.operator.scalar.CallOperator;
 import com.starrocks.sql.optimizer.operator.scalar.CastOperator;
 import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
 import com.starrocks.sql.optimizer.operator.scalar.CompoundPredicateOperator;
@@ -20,6 +29,10 @@
 import com.starrocks.sql.optimizer.rewrite.scalar.SimplifiedPredicateRule;
 import org.junit.Test;
 
+import java.time.LocalDateTime;
+import java.util.Arrays;
+import java.util.stream.Collectors;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -119,4 +132,32 @@ public void testNormalizeIsNull() {
                 .rewrite(isnotNull, ScalarOperatorRewriter.DEFAULT_REWRITE_SCAN_PREDICATE_RULES);
         assertEquals(ConstantOperator.TRUE, rewritten2);
     }
+
+    @Test
+    public void testRangeExtract() {
+        Supplier<ScalarOperator> predicate1Maker = () -> {
+            ColumnRefOperator col1 = new ColumnRefOperator(1, Type.DATE, "dt", false);
+            CallOperator call = new CallOperator(FunctionSet.DATE_TRUNC, Type.DATE,
+                    Arrays.asList(ConstantOperator.createVarchar("YEAR"), col1));
+            return new BinaryPredicateOperator(BinaryType.EQ, call, ConstantOperator.createDate(
+                    LocalDateTime.of(2024, 1, 1, 0, 0, 0)));
+        };
+
+        Supplier<ScalarOperator> predicate2Maker = () -> {
+            ColumnRefOperator col2 = new ColumnRefOperator(2, Type.VARCHAR, "mode", false);
+            return new BinaryPredicateOperator(BinaryType.EQ, col2, ConstantOperator.createVarchar("Buzz"));
+        };
+
+        ScalarOperator predicate1 = Utils.compoundAnd(predicate1Maker.get(), predicate2Maker.get());
+        ScalarOperator predicate2 = Utils.compoundAnd(predicate1Maker.get(), predicate2Maker.get());
+        ScalarOperator predicates = Utils.compoundAnd(predicate1, predicate2);
+
+        ScalarRangePredicateExtractor rangeExtractor = new ScalarRangePredicateExtractor();
+        ScalarOperator result = rangeExtractor.rewriteOnlyColumn(Utils.compoundAnd(Utils.extractConjuncts(predicates)
+                .stream().map(rangeExtractor::rewriteOnlyColumn).collect(Collectors.toList())));
+        Preconditions.checkState(result != null);
+        String expect = "date_trunc(YEAR, 1: dt) = 2024-01-01 AND 2: mode = Buzz";
+        String actual = result.toString();
+        Assert.assertEquals(actual, expect, actual);
+    }
 }