[BugFix] Add implicit cast rule for list partitino pruner with generated column (backport StarRocks#54543) (StarRocks#57697)

Co-authored-by: shuming.li <[email protected]>

Author: mergify[bot]

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ListPartitionPruner.java

@@ -438,8 +438,10 @@ public ScalarOperator visitInPredicate(InPredicateOperator operator, Void contex
             return null;
         }
 
-        // Fold constants
         ScalarOperatorRewriter rewriter = new ScalarOperatorRewriter();
+        // implicit cast
+        result = rewriter.rewrite(result, ScalarOperatorRewriter.DEFAULT_TYPE_CAST_RULE);
+        // fold constant
         result = rewriter.rewrite(result, ScalarOperatorRewriter.FOLD_CONSTANT_RULES);
         return result;
     }

fe/fe-core/src/test/java/com/starrocks/sql/plan/PartitionPruneTest.java

@@ -25,6 +25,7 @@
 import com.starrocks.sql.optimizer.operator.logical.LogicalScanOperator;
 import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 import com.starrocks.sql.optimizer.statistics.StatisticsCalculator;
+import com.starrocks.thrift.TExplainLevel;
 import com.starrocks.utframe.UtFrameUtils;
 import org.junit.Assert;
 import org.junit.BeforeClass;
@@ -75,6 +76,22 @@ public static void beforeClass() throws Exception {
         starRocksAssert.ddl("ALTER TABLE t_gen_col ADD PARTITION p2_202402 VALUES IN (('2', '2024-02-01'))");
         starRocksAssert.ddl("ALTER TABLE t_gen_col ADD PARTITION p2_202403 VALUES IN (('2', '2024-03-01'))");
 
+        // date_trunc('month', hours_add(date_trunc('day', hours_sub(c1, 8)), 8))
+        starRocksAssert.withTable("CREATE TABLE t_gen_col2 (" +
+                " c1 datetime NOT NULL," +
+                " c2 bigint," +
+                " c3 DATETIME NULL AS date_trunc('month', hours_add(date_trunc('day', hours_sub(c1, 8)), 8)) " +
+                " ) " +
+                " DUPLICATE KEY(c1) " +
+                " PARTITION BY (c2, c3) " +
+                " PROPERTIES('replication_num'='1')");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p1_202401 VALUES IN (('1', '2024-01-01'))");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p1_202402 VALUES IN (('1', '2024-02-01'))");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p1_202403 VALUES IN (('1', '2024-03-01'))");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p2_202401 VALUES IN (('2', '2024-01-01'))");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p2_202402 VALUES IN (('2', '2024-02-01'))");
+        starRocksAssert.ddl("ALTER TABLE t_gen_col2 ADD PARTITION p2_202403 VALUES IN (('2', '2024-03-01'))");
+
         starRocksAssert.withTable("CREATE TABLE t_bool_partition (" +
                 " c1 datetime NOT NULL, " +
                 " c2 boolean" +
@@ -224,6 +241,7 @@ public void testNullException() throws Exception {
     private static Pair<ScalarOperator, LogicalScanOperator> buildConjunctAndScan(String sql) throws Exception {
         Pair<String, ExecPlan> pair = UtFrameUtils.getPlanAndFragment(connectContext, sql);
         ExecPlan execPlan = pair.second;
+        System.out.println(execPlan.getExplainString(TExplainLevel.NORMAL));
         LogicalScanOperator scanOperator =
                 (LogicalScanOperator) execPlan.getLogicalPlan().getRoot().inputAt(0).inputAt(0).inputAt(0).getOp();
         ScalarOperator predicate = execPlan.getPhysicalPlan().getOp().getPredicate();
@@ -238,9 +256,15 @@ private void testRemovePredicate(String sql, String expected) throws Exception {
         Assert.assertEquals(expected, newPredicate.toString());
     }
 
+    private void testAssertContains(String sql, String expected) throws Exception {
+        Pair<String, ExecPlan> pair = UtFrameUtils.getPlanAndFragment(connectContext, sql);
+        ExecPlan execPlan = pair.second;
+        String plan = execPlan.getExplainString(TExplainLevel.NORMAL);
+        PlanTestBase.assertContains(plan, expected);
+    }
+
     @Test
     public void testGeneratedColumnPrune_RemovePredicate() throws Exception {
-        testRemovePredicate("select * from t_gen_col where c1 = '2024-01-01' ", "true");
         testRemovePredicate("select * from t_gen_col where c1 = '2024-01-01' and c2 > 100", "true");
         testRemovePredicate("select * from t_gen_col where c1 >= '2024-01-01'  and c1 <= '2024-01-03' " +
                 "and c2 > 100", "true");
@@ -258,6 +282,23 @@ public void testGeneratedColumnPrune_RemovePredicate() throws Exception {
                 "cast(add(2: c2, 100) as double) > add(cast(1: c1 as double), 1)");
     }
 
+    @Test
+    public void testGeneratedColumnPrune_RemovePredicate2() throws Exception {
+        testAssertContains("select * from t_gen_col2 where c1 >= '2024-02-02' ", "partitions=4/6");
+        testAssertContains("select * from t_gen_col2 where c1 = '2024-02-02' ", "partitions=2/6");
+        testAssertContains("select * from t_gen_col2 where c1 = '2024-02-02' and c2 > 100", "partitions=0/6");
+        testAssertContains("select * from t_gen_col2 where c1 >= '2024-02-02'  and c1 <= '2024-02-03' " +
+                "and c2 > 100", "partitions=0/6");
+        testAssertContains("select * from t_gen_col2 where c2 in (1, 2,3)", "partitions=6/6");
+        testAssertContains("select * from t_gen_col2 where c2 = cast('123' as int)", "partitions=0/6");
+
+        // can not be removed
+        testAssertContains("select * from t_gen_col2 where c1 = random() and c2 > 100",
+                "partitions=0/6");
+        testAssertContains("select * from t_gen_col2 where c2 + 100 > c1 + 1",
+                "partitions=6/6");
+    }
+
     @Test
     public void testGeneratedColumnPrune() throws Exception {
         // c2