branch-3.1: [test](nereids) add ut for sql cache/sorted partition cache #55520

fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSortedPartitionsCacheManager.java

@@ -35,6 +35,7 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Range;
 import lombok.AllArgsConstructor;
 import lombok.Data;
@@ -105,6 +106,11 @@ public Optional<SortedPartitionRanges<?>> get(
         return Optional.of(partitionCacheContext.sortedPartitionRanges);
     }
 
+    @VisibleForTesting
+    public Cache<TableIdentifier, PartitionCacheContext> getPartitionCaches() {
+        return partitionCaches;
+    }
+
     private SortedPartitionRanges<?> loadCache(
             TableIdentifier key, SupportBinarySearchFilteringPartitions table, CatalogRelation scan)
             throws RpcException {
@@ -182,6 +188,7 @@ public static synchronized void updateConfig() {
                 Config.expire_cache_partition_meta_table_in_fe_second
         );
         caches.putAll(cacheManager.partitionCaches.asMap());
+        caches.cleanUp();
         cacheManager.partitionCaches = caches;
     }
 

fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSqlCacheManager.java

@@ -67,6 +67,7 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import org.apache.commons.collections.CollectionUtils;
@@ -98,6 +99,11 @@ public NereidsSqlCacheManager() {
         );
     }
 
+    @VisibleForTesting
+    public Cache<String, SqlCacheContext> getSqlCaches() {
+        return sqlCaches;
+    }
+
     public static synchronized void updateConfig() {
         Env currentEnv = Env.getCurrentEnv();
         if (currentEnv == null) {
@@ -113,6 +119,7 @@ public static synchronized void updateConfig() {
                 Config.expire_sql_cache_in_fe_second
         );
         sqlCaches.putAll(sqlCacheManager.sqlCaches.asMap());
+        sqlCaches.cleanUp();
         sqlCacheManager.sqlCaches = sqlCaches;
     }
 

fe/fe-core/src/test/java/org/apache/doris/qe/SortedPartitionsCacheTest.java

@@ -0,0 +1,81 @@
+// 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.qe;
+
+import org.apache.doris.catalog.Env;
+import org.apache.doris.common.cache.NereidsSortedPartitionsCacheManager;
+import org.apache.doris.utframe.TestWithFeService;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SortedPartitionsCacheTest extends TestWithFeService {
+
+    @Test
+    public void testPartitionCache() throws Exception {
+        createDatabase("test");
+        createTable("create table test.key_1_fixed_range_date_part (a int, dt datetime, c varchar(100)) duplicate key(a)\n"
+                + "partition by range(dt) (\n"
+                + "    PARTITION p_min VALUES LESS THAN (\"2023-01-01 00:00:00\"),\n"
+                + "    PARTITION p_202301 VALUES [('2023-01-01 00:00:00'), ('2023-02-01 00:00:00')),\n"
+                + "    PARTITION p_202302 VALUES [('2023-02-01 00:00:00'), ('2023-03-01 00:00:00')),\n"
+                + "    PARTITION p_202303 VALUES [('2023-03-01 00:00:00'), ('2023-04-01 00:00:00')),\n"
+                + "    PARTITION p_202304 VALUES [('2023-04-01 00:00:00'), ('2023-05-01 00:00:00')),\n"
+                + "    PARTITION p_202305 VALUES [('2023-05-01 00:00:00'), ('2023-06-01 00:00:00')),\n"
+                + "    PARTITION p_202306 VALUES [('2023-06-01 00:00:00'), ('2023-07-01 00:00:00')),\n"
+                + "    PARTITION p_202307 VALUES [('2023-07-01 00:00:00'), ('2023-08-01 00:00:00')),\n"
+                + "    PARTITION p_202308 VALUES [('2023-08-01 00:00:00'), ('2023-09-01 00:00:00')),\n"
+                + "    PARTITION p_202309 VALUES [('2023-09-01 00:00:00'), ('2023-10-01 00:00:00')),\n"
+                + "    PARTITION p_202310 VALUES [('2023-10-01 00:00:00'), ('2023-11-01 00:00:00')),\n"
+                + "    PARTITION p_202311 VALUES [('2023-11-01 00:00:00'), ('2023-12-01 00:00:00')),\n"
+                + "    PARTITION p_202312 VALUES [('2023-12-01 00:00:00'), ('2024-01-01 00:00:00')),\n"
+                + "    PARTITION p_max VALUES [('2024-01-01 00:00:00'), ('9999-12-31 23:59:59'))\n"
+                + ") distributed by hash(a) properties(\"replication_num\"=\"1\");");
+
+        createTable("create table test.key_1_fixed_range_date_part2 (a int, dt datetime, c varchar(100)) duplicate key(a)\n"
+                + "partition by range(dt) (\n"
+                + "    PARTITION p_min VALUES LESS THAN (\"2023-01-01 00:00:00\"),\n"
+                + "    PARTITION p_202301 VALUES [('2023-01-01 00:00:00'), ('2023-02-01 00:00:00')),\n"
+                + "    PARTITION p_202302 VALUES [('2023-02-01 00:00:00'), ('2023-03-01 00:00:00')),\n"
+                + "    PARTITION p_202303 VALUES [('2023-03-01 00:00:00'), ('2023-04-01 00:00:00')),\n"
+                + "    PARTITION p_202304 VALUES [('2023-04-01 00:00:00'), ('2023-05-01 00:00:00')),\n"
+                + "    PARTITION p_202305 VALUES [('2023-05-01 00:00:00'), ('2023-06-01 00:00:00')),\n"
+                + "    PARTITION p_202306 VALUES [('2023-06-01 00:00:00'), ('2023-07-01 00:00:00')),\n"
+                + "    PARTITION p_202307 VALUES [('2023-07-01 00:00:00'), ('2023-08-01 00:00:00')),\n"
+                + "    PARTITION p_202308 VALUES [('2023-08-01 00:00:00'), ('2023-09-01 00:00:00')),\n"
+                + "    PARTITION p_202309 VALUES [('2023-09-01 00:00:00'), ('2023-10-01 00:00:00')),\n"
+                + "    PARTITION p_202310 VALUES [('2023-10-01 00:00:00'), ('2023-11-01 00:00:00')),\n"
+                + "    PARTITION p_202311 VALUES [('2023-11-01 00:00:00'), ('2023-12-01 00:00:00')),\n"
+                + "    PARTITION p_202312 VALUES [('2023-12-01 00:00:00'), ('2024-01-01 00:00:00')),\n"
+                + "    PARTITION p_max VALUES [('2024-01-01 00:00:00'), ('9999-12-31 23:59:59'))\n"
+                + ") distributed by hash(a) properties(\"replication_num\"=\"1\");");
+
+        // wait cache enable
+        Thread.sleep(10 * 1000);
+
+        executeNereidsSql("select * from test.key_1_fixed_range_date_part where dt='2023-07-01 05:00:00'");
+        executeNereidsSql("select * from test.key_1_fixed_range_date_part2 where dt='2023-07-01 05:00:00'");
+
+        Env currentEnv = Env.getCurrentEnv();
+        NereidsSortedPartitionsCacheManager sortedPartitionsCacheManager = currentEnv.getSortedPartitionsCacheManager();
+        Assertions.assertEquals(2, sortedPartitionsCacheManager.getPartitionCaches().asMap().size());
+
+        executeNereidsSql("admin set frontend config ('cache_partition_meta_table_manage_num'='1')");
+        Assertions.assertEquals(1, sortedPartitionsCacheManager.getPartitionCaches().asMap().size());
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/qe/SqlCacheTest.java

@@ -18,17 +18,20 @@
 package org.apache.doris.qe;
 
 import org.apache.doris.analysis.UserIdentity;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.common.cache.NereidsSqlCacheManager;
 import org.apache.doris.nereids.SqlCacheContext;
 import org.apache.doris.proto.Types.PUniqueId;
 import org.apache.doris.thrift.TUniqueId;
+import org.apache.doris.utframe.TestWithFeService;
 
 import com.google.common.collect.ImmutableSet;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import java.util.UUID;
 
-public class SqlCacheTest {
+public class SqlCacheTest extends TestWithFeService {
     @Test
     public void testCacheKey() {
         TUniqueId queryId = new TUniqueId();
@@ -57,4 +60,18 @@ public void testCacheKey() {
         PUniqueId key3 = cacheContext3.doComputeCacheKeyMd5(ImmutableSet.of());
         Assertions.assertNotEquals(key1, key3);
     }
+
+    @Test
+    public void testSqlCache() throws Exception {
+        connectContext.getSessionVariable().setEnableSqlCache(true);
+        executeNereidsSql("select 100");
+        executeNereidsSql("select 200");
+
+        Env currentEnv = Env.getCurrentEnv();
+        NereidsSqlCacheManager sqlCacheManager = currentEnv.getSqlCacheManager();
+        Assertions.assertEquals(2, sqlCacheManager.getSqlCaches().asMap().size());
+
+        executeNereidsSql("admin set frontend config ('sql_cache_manage_num'='1')");
+        Assertions.assertEquals(1, sqlCacheManager.getSqlCaches().asMap().size());
+    }
 }

fe/fe-core/src/test/java/org/apache/doris/utframe/TestWithFeService.java

@@ -626,6 +626,24 @@ public void executeSql(String queryStr) throws Exception {
         }
     }
 
+    public void executeNereidsSql(String queryStr) throws Exception {
+        connectContext.getState().reset();
+
+        StatementContext statementContext = new StatementContext(connectContext, new OriginStatement(queryStr, 0));
+        connectContext.setStatementContext(statementContext);
+        statementContext.setConnectContext(connectContext);
+
+        LogicalPlan plan = new NereidsParser().parseSingle(queryStr);
+        LogicalPlanAdapter logicalPlanAdapter = new LogicalPlanAdapter(plan, statementContext);
+        logicalPlanAdapter.setOrigStmt(statementContext.getOriginStatement());
+        StmtExecutor stmtExecutor = new StmtExecutor(connectContext, logicalPlanAdapter);
+        stmtExecutor.execute();
+        if (connectContext.getState().getStateType() == QueryState.MysqlStateType.ERR
+                || connectContext.getState().getErrorCode() != null) {
+            throw new IllegalStateException(connectContext.getState().getErrorMessage());
+        }
+    }
+
     public void createDatabase(String db) throws Exception {
         String createDbStmtStr = "CREATE DATABASE " + db;
         CreateDbStmt createDbStmt = (CreateDbStmt) parseAndAnalyzeStmt(createDbStmtStr);