branch-4.1: [fix](point query) Keep LogicalOlapScan for short-circuit point query on empty table #62948

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),

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<Long> partitionIdsToPrune = scan.getSelectedPartitionIds();

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;
+        }
+    }
+}

regression-test/data/point_query_p0/test_point_query.out

@@ -152,6 +152,8 @@
 -- !sql --
 0	1111111
 
+-- !sql_empty --
+
 -- !sql --
 10	20	aabc	value
 

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');"