[BugFix] fix hive partition prune when casting string to int on in_pr…

…edicate (backport #44875) (#44891) Co-authored-by: stephen <[email protected]>
StarRocks · Apr 28, 2024 · ddd7b4f · ddd7b4f
1 parent d4a79f8
commit ddd7b4f
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/...re/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ListPartitionPruner.java b/...re/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ListPartitionPruner.java
@@ -417,6 +417,13 @@ private Set<Long> evalInPredicate(InPredicateOperator inPredicate) {
         Set<Long> matches = Sets.newHashSet();
         ConcurrentNavigableMap<LiteralExpr, Set<Long>> partitionValueMap = columnToPartitionValuesMap.get(child);
         Set<Long> nullPartitions = columnToNullPartitions.get(child);
+
+        if (inPredicate.getChild(0) instanceof CastOperator && partitionValueMap != null) {
+            // partitionValueMap need cast to target type
+            partitionValueMap = getCastPartitionValueMap((CastOperator) inPredicate.getChild(0),
+                    partitionValueMap);
+        }
+
         if (partitionValueMap == null || nullPartitions == null || partitionValueMap.isEmpty()) {
             return null;
         }

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/optimizer/rule/ListPartitionPrunerTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/optimizer/rule/ListPartitionPrunerTest.java
@@ -442,6 +442,36 @@ public void testExternalTableBinaryPredicateWithException() throws AnalysisExcep
         Assert.assertNull(pruner.prune());
     }
 
+    @Test
+    public void testStringCastInPredicate() throws AnalysisException {
+        // string_col="1"   1
+        // string_col="2"  2
+        // string_col="3"   3
+
+        ColumnRefOperator stringColumn = new ColumnRefOperator(1, Type.STRING, "string_col", true);
+
+        // column -> partition values
+        Map<ColumnRefOperator, ConcurrentNavigableMap<LiteralExpr, Set<Long>>> columnToPartitionValuesMap =
+                Maps.newConcurrentMap();
+        ConcurrentNavigableMap<LiteralExpr, Set<Long>> stringPartitionValuesMap = new ConcurrentSkipListMap<>();
+        columnToPartitionValuesMap.put(stringColumn, stringPartitionValuesMap);
+        stringPartitionValuesMap.put(new StringLiteral("1"), Sets.newHashSet(1L));
+        stringPartitionValuesMap.put(new StringLiteral("2"), Sets.newHashSet(2L));
+        stringPartitionValuesMap.put(new StringLiteral("3"), Sets.newHashSet(3L));
+
+        Map<ColumnRefOperator, Set<Long>> columnToNullPartitions = Maps.newHashMap();
+        columnToNullPartitions.put(stringColumn, Sets.newHashSet(9L));
+
+        List<ScalarOperator> conjuncts = Lists.newArrayList();
+        ListPartitionPruner pruner =
+                new ListPartitionPruner(columnToPartitionValuesMap, columnToNullPartitions, conjuncts, null);
+        conjuncts.clear();
+        conjuncts.add(new InPredicateOperator(new CastOperator(Type.INT, stringColumn),
+                ConstantOperator.createInt(2),
+                ConstantOperator.createInt(3)));
+        Assert.assertEquals(2, pruner.prune().size());
+    }
+
     @Test
     public void testSpecifyPartition() throws AnalysisException {
         // 2 partition columns