diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java index 3947f07509bf38..1e1ecc29c86d37 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java @@ -722,6 +722,10 @@ public class Rewriter extends AbstractBatchJobExecutor { topic("Table/Physical optimization", cascadesContext -> cascadesContext.rewritePlanContainsTypes(LogicalCatalogRelation.class), topDown( + // Mark short-circuit point query before pruning empty partitions, + // otherwise PRUNE_EMPTY_PARTITION may replace LogicalOlapScan with LogicalEmptyRelation + // and short-circuit flag can never be set. + new LogicalResultSinkToShortCircuitPointQuery(), new PruneOlapScanPartition(), new PruneEmptyPartition(), new PruneFileScanPartition(), @@ -734,8 +738,6 @@ public class Rewriter extends AbstractBatchJobExecutor { topic("adjust preagg status", custom(RuleType.SET_PREAGG_STATUS, SetPreAggStatus::new) ), - topic("Point query short circuit", - topDown(new LogicalResultSinkToShortCircuitPointQuery())), topic("eliminate", // SORT_PRUNING should be applied after mergeLimit custom(RuleType.ELIMINATE_SORT, EliminateSort::new), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java index 3347518b1652a8..aadd48bbd523be 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java @@ -40,6 +40,12 @@ public class PruneEmptyPartition extends OneRewriteRuleFactory { @Override public Rule build() { return logicalOlapScan().thenApply(ctx -> { + // We still want to keep LogicalOlapScan even if partitions are empty, + // so that the planner can build a scan node and the PreparedStatement can cache ShortCircuitQueryContext. + if (ctx.connectContext != null && ctx.connectContext.getStatementContext() != null + && ctx.connectContext.getStatementContext().isShortCircuitQuery()) { + return null; + } LogicalOlapScan scan = ctx.root; OlapTable table = scan.getTable(); List partitionIdsToPrune = scan.getSelectedPartitionIds(); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java new file mode 100644 index 00000000000000..e036b3a7c49b21 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java @@ -0,0 +1,78 @@ +// 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.rewrite; + +import org.apache.doris.common.FeConstants; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.util.MemoPatternMatchSupported; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Regression test: + * For short-circuit point query, we should not rewrite LogicalOlapScan to LogicalEmptyRelation + * even if the table partitions are empty. Otherwise PreparedStatement could not cache + * ShortCircuitQueryContext and may downgrade to normal planning repeatedly. + */ +class ShortCircuitPointQueryTest extends TestWithFeService + implements MemoPatternMatchSupported { + + @Override + protected void runBeforeAll() throws Exception { + createDatabase("test"); + useDatabase("test"); + createTable("CREATE TABLE `tbl_point_query` (\n" + + " `key` int(11) NULL,\n" + + " `v1` varchar(30) NULL\n" + + ") ENGINE=OLAP\n" + + "UNIQUE KEY(`key`)\n" + + "DISTRIBUTED BY HASH(`key`) BUCKETS 1\n" + + "PROPERTIES (\n" + + " \"replication_num\" = \"1\",\n" + + " \"enable_unique_key_merge_on_write\" = \"true\",\n" + + " \"light_schema_change\" = \"true\",\n" + + " \"store_row_column\" = \"true\"\n" + + ");"); + } + + @Test + void testShortCircuitPointQueryKeepOlapScanWhenTableEmpty() { + boolean originRunningUnitTest = FeConstants.runningUnitTest; + FeConstants.runningUnitTest = false; + try { + String sql = "select * from tbl_point_query where `key` = 1"; + Plan plan = PlanChecker.from(connectContext) + .analyze(sql) + .rewrite() + .getPlan(); + + // short-circuit flag should be set + Assertions.assertTrue(connectContext.getStatementContext().isShortCircuitQuery()); + // should still keep scan node for point query path + Assertions.assertTrue(plan.anyMatch(p -> p instanceof LogicalOlapScan)); + Assertions.assertFalse(plan.anyMatch(p -> p instanceof LogicalEmptyRelation)); + } finally { + FeConstants.runningUnitTest = originRunningUnitTest; + } + } +} diff --git a/regression-test/data/point_query_p0/test_point_query.out b/regression-test/data/point_query_p0/test_point_query.out index 5fafa2fb91db33..e03c07418650e0 100644 --- a/regression-test/data/point_query_p0/test_point_query.out +++ b/regression-test/data/point_query_p0/test_point_query.out @@ -152,6 +152,8 @@ -- !sql -- 0 1111111 +-- !sql_empty -- + -- !sql -- 10 20 aabc value diff --git a/regression-test/suites/point_query_p0/test_point_query.groovy b/regression-test/suites/point_query_p0/test_point_query.groovy index 90c6bacf599ee4..431687c63a783d 100644 --- a/regression-test/suites/point_query_p0/test_point_query.groovy +++ b/regression-test/suites/point_query_p0/test_point_query.groovy @@ -328,6 +328,11 @@ suite("test_point_query") { DISTRIBUTED BY HASH(`col1`, `col2`, `loc3`) BUCKETS 1 PROPERTIES ( "replication_allocation" = "tag.location.default: 1", "bloom_filter_columns" = "col1", "store_row_column" = "true", "enable_mow_light_delete" = "false" ); """ + explain { + sql("select * from table_3821461 where col1 = -10 and col2 = 20 and loc3 = 'aabc'") + contains "SHORT-CIRCUIT" + } + qt_sql_empty "select * from table_3821461 where col1 = 10 and col2 = 20 and loc3 = 'aabc';" sql "insert into table_3821461 values (-10, 20, 'aabc', 'value')" sql "insert into table_3821461 values (10, 20, 'aabc', 'value');" sql "insert into table_3821461 values (20, 30, 'aabc', 'value');"