[BugFix] Fix cancel query for external low cardinality (backport StarRocks#58278) (backport StarRocks#58342) (StarRocks#59068)

mergify[bot] · ZiheLiu · web-flow · commit 1c6d046edd22 · 2025-05-19T10:17:23.000Z
Signed-off-by: zihe.liu &lt;ziheliu1024@gmail.com&gt;
Co-authored-by: mergify[bot] &lt;37929162+mergify[bot]@users.noreply.github.com&gt;
Co-authored-by: zihe.liu &lt;ziheliu1024@gmail.com&gt;
diff --git a/be/src/exec/pipeline/fragment_context.cpp b/be/src/exec/pipeline/fragment_context.cpp
@@ -244,9 +244,12 @@ void FragmentContext::set_stream_load_contexts(const std::vector<StreamLoadConte
 }
 
 // Note: this function should be thread safe
-void FragmentContext::cancel(const Status& status) {
+void FragmentContext::cancel(const Status& status, bool cancelled_by_fe) {
     if (!status.ok() && _runtime_state != nullptr && _runtime_state->query_ctx() != nullptr) {
         _runtime_state->query_ctx()->release_workgroup_token_once();
+        if (cancelled_by_fe) {
+            _runtime_state->query_ctx()->set_cancelled_by_fe();
+        }
     }
 
     _runtime_state->set_is_cancelled(true);
diff --git a/be/src/exec/pipeline/fragment_context.h b/be/src/exec/pipeline/fragment_context.h
@@ -90,7 +90,7 @@ class FragmentContext {
         return status == nullptr ? Status::OK() : *status;
     }
 
-    void cancel(const Status& status);
+    void cancel(const Status& status, bool cancelled_by_fe = false);
 
     void finish() { cancel(Status::OK()); }
 
diff --git a/be/src/exec/pipeline/fragment_executor.cpp b/be/src/exec/pipeline/fragment_executor.cpp
@@ -99,6 +99,7 @@ Status FragmentExecutor::_prepare_query_ctx(ExecEnv* exec_env, const UnifiedExec
     const auto& query_id = params.query_id;
     const auto& fragment_instance_id = request.fragment_instance_id();
     const auto& query_options = request.common().query_options;
+    const auto& t_desc_tbl = request.common().desc_tbl;
 
     auto&& existing_query_ctx = exec_env->query_context_mgr()->get(query_id);
     if (existing_query_ctx) {
@@ -108,7 +109,8 @@ Status FragmentExecutor::_prepare_query_ctx(ExecEnv* exec_env, const UnifiedExec
         }
     }
 
-    ASSIGN_OR_RETURN(_query_ctx, exec_env->query_context_mgr()->get_or_register(query_id));
+    const bool query_ctx_should_exist = t_desc_tbl.__isset.is_cached && t_desc_tbl.is_cached;
+    ASSIGN_OR_RETURN(_query_ctx, exec_env->query_context_mgr()->get_or_register(query_id, query_ctx_should_exist));
     _query_ctx->set_exec_env(exec_env);
     if (params.__isset.instances_number) {
         _query_ctx->set_total_fragments(params.instances_number);
diff --git a/be/src/exec/pipeline/query_context.cpp b/be/src/exec/pipeline/query_context.cpp
@@ -110,7 +110,7 @@ void QueryContext::cancel(const Status& status, bool cancelled_by_fe) {
     _is_cancelled = true;
     if (cancelled_by_fe) {
         // only update when confirm cancelled from fe
-        _cancelled_by_fe = true;
+        set_cancelled_by_fe();
     }
     if (_cancelled_status.load() != nullptr) {
         return;
@@ -387,7 +387,8 @@ QueryContextManager::~QueryContextManager() {
         return query_ctx->get_cancelled_status();           \
     }
 
-StatusOr<QueryContext*> QueryContextManager::get_or_register(const TUniqueId& query_id) {
+StatusOr<QueryContext*> QueryContextManager::get_or_register(const TUniqueId& query_id,
+                                                             bool return_error_if_not_exist) {
     size_t i = _slot_idx(query_id);
     auto& mutex = _mutexes[i];
     auto& context_map = _context_maps[i];
@@ -422,6 +423,10 @@ StatusOr<QueryContext*> QueryContextManager::get_or_register(const TUniqueId& qu
             }
         }
 
+        if (return_error_if_not_exist) {
+            return Status::Cancelled("Query terminates prematurely");
+        }
+
         // finally, find no query contexts, so create a new one
         auto&& ctx = std::make_shared<QueryContext>();
         auto* ctx_raw_ptr = ctx.get();
diff --git a/be/src/exec/pipeline/query_context.h b/be/src/exec/pipeline/query_context.h
@@ -154,6 +154,7 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
     FragmentContextManager* fragment_mgr();
 
     void cancel(const Status& status, bool cancelled_by_fe);
+    void set_cancelled_by_fe() { _cancelled_by_fe = true; }
 
     void set_is_runtime_filter_coordinator(bool flag) { _is_runtime_filter_coordinator = flag; }
 
@@ -386,7 +387,7 @@ class QueryContextManager {
     QueryContextManager(size_t log2_num_slots);
     ~QueryContextManager();
     Status init();
-    StatusOr<QueryContext*> get_or_register(const TUniqueId& query_id);
+    StatusOr<QueryContext*> get_or_register(const TUniqueId& query_id, bool return_error_if_not_exist = false);
     QueryContextPtr get(const TUniqueId& query_id, bool need_prepared = false);
     size_t size();
     bool remove(const TUniqueId& query_id);
diff --git a/be/src/formats/parquet/column_reader_factory.cpp b/be/src/formats/parquet/column_reader_factory.cpp
@@ -18,9 +18,12 @@
 #include "formats/parquet/scalar_column_reader.h"
 #include "formats/parquet/schema.h"
 #include "formats/utils.h"
+#include "util/failpoint/fail_point.h"
 
 namespace starrocks::parquet {
 
+DEFINE_FAIL_POINT(parquet_reader_returns_global_dict_not_match_status);
+
 StatusOr<ColumnReaderPtr> ColumnReaderFactory::create(const ColumnReaderOptions& opts, const ParquetField* field,
                                                       const TypeDescriptor& col_type) {
     // We will only set a complex type in ParquetField
@@ -158,6 +161,11 @@ StatusOr<ColumnReaderPtr> ColumnReaderFactory::create(const ColumnReaderOptions&
 
 StatusOr<ColumnReaderPtr> ColumnReaderFactory::create(ColumnReaderPtr ori_reader, const GlobalDictMap* dict,
                                                       SlotId slot_id, int64_t num_rows) {
+    FAIL_POINT_TRIGGER_EXECUTE(parquet_reader_returns_global_dict_not_match_status, {
+        return Status::GlobalDictNotMatch(
+                fmt::format("SlotId: {}, Not dict encoded and not low rows on global dict column. ", slot_id));
+    });
+
     if (ori_reader->get_column_parquet_field()->type == ColumnType::ARRAY) {
         ASSIGN_OR_RETURN(ColumnReaderPtr child_reader,
                          ColumnReaderFactory::create(
diff --git a/be/src/service/internal_service.cpp b/be/src/service/internal_service.cpp
@@ -564,7 +564,7 @@ void PInternalServiceImplBase<T>::_cancel_plan_fragment(google::protobuf::RpcCon
                         "FragmentContext already destroyed: query_id=$0, fragment_instance_id=$1", print_id(query_id),
                         print_id(tid));
             } else {
-                fragment_ctx->cancel(Status::Cancelled(reason_string));
+                fragment_ctx->cancel(Status::Cancelled(reason_string), true);
             }
         }
     } else {
diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteExceptionHandler.java b/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteExceptionHandler.java
@@ -38,6 +38,7 @@
 import com.starrocks.sql.optimizer.statistics.IRelaxDictManager;
 import com.starrocks.sql.plan.ExecPlan;
 import com.starrocks.thrift.TExplainLevel;
+import com.starrocks.thrift.TStatusCode;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -66,6 +67,12 @@ public static void handle(Exception e, RetryContext context) throws Exception {
         }
     }
 
+    public static boolean isRetryableStatus(TStatusCode statusCode) {
+        return statusCode == TStatusCode.REMOTE_FILE_NOT_FOUND
+                || statusCode == TStatusCode.THRIFT_RPC_ERROR
+                || statusCode == TStatusCode.GLOBAL_DICT_NOT_MATCH;
+    }
+
     // If modifications are made to the partition files of a Hive table by user,
     // such as through "insert overwrite partition", the Frontend couldn't be aware of these changes.
     // As a result, queries may use the file information cached in the FE for execution.
diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/scheduler/Deployer.java b/fe/fe-core/src/main/java/com/starrocks/qe/scheduler/Deployer.java
@@ -23,6 +23,7 @@
 import com.starrocks.common.profile.Timer;
 import com.starrocks.common.profile.Tracers;
 import com.starrocks.qe.ConnectContext;
+import com.starrocks.qe.ExecuteExceptionHandler;
 import com.starrocks.qe.scheduler.dag.ExecutionDAG;
 import com.starrocks.qe.scheduler.dag.ExecutionFragment;
 import com.starrocks.qe.scheduler.dag.FragmentInstance;
@@ -262,6 +263,12 @@ private void waitForDeploymentCompletion(List<FragmentInstanceExecState> executi
             if (firstErrResult == null) {
                 firstErrResult = res;
                 firstErrExecution = execution;
+            } else if (firstErrResult.getStatusCode() == TStatusCode.CANCELLED &&
+                    ExecuteExceptionHandler.isRetryableStatus(res.getStatusCode())) {
+                // If the first error is cancelled and the subsequent error is retryable, we store the latter to give a chance
+                // to retry this query.
+                firstErrResult = res;
+                firstErrExecution = execution;
             }
             if (TStatusCode.TIMEOUT == res.getStatusCode()) {
                 break;
diff --git a/fe/fe-core/src/test/java/com/starrocks/qe/scheduler/StartSchedulingTest.java b/fe/fe-core/src/test/java/com/starrocks/qe/scheduler/StartSchedulingTest.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.Maps;
 import com.starrocks.common.Reference;
 import com.starrocks.common.StarRocksException;
+import com.starrocks.connector.exception.GlobalDictNotMatchException;
 import com.starrocks.proto.PCancelPlanFragmentRequest;
 import com.starrocks.proto.PCancelPlanFragmentResult;
 import com.starrocks.proto.PExecPlanFragmentResult;
@@ -54,6 +55,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static com.starrocks.utframe.MockedBackend.MockPBackendService;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -212,6 +214,57 @@ public Future<PExecPlanFragmentResult> execPlanFragmentAsync(PExecPlanFragmentRe
         Assert.assertTrue(scheduler.isDone());
     }
 
+    @Test
+    public void testDeployFutureThrowRetryableException() throws Exception {
+        AtomicInteger numDeployedFragments = new AtomicInteger(0);
+
+        setBackendService(address -> {
+
+            final int numFragments = numDeployedFragments.incrementAndGet();
+            if (numFragments == 1) {
+                return new MockPBackendService();
+            }
+
+            if (numFragments == 2) {
+                return new MockPBackendService() {
+                    @Override
+                    public Future<PExecPlanFragmentResult> execPlanFragmentAsync(PExecPlanFragmentRequest request) {
+                        return submit(() -> {
+                            PExecPlanFragmentResult result = new PExecPlanFragmentResult();
+                            StatusPB pStatus = new StatusPB();
+                            pStatus.statusCode = TStatusCode.CANCELLED.getValue();
+                            pStatus.errorMsgs = Collections.singletonList("test CANCELLED error message");
+                            result.status = pStatus;
+                            return result;
+                        });
+                    }
+                };
+            }
+
+            if (numFragments == 3) {
+                return new MockPBackendService() {
+                    @Override
+                    public Future<PExecPlanFragmentResult> execPlanFragmentAsync(PExecPlanFragmentRequest request) {
+                        return submit(() -> {
+                            PExecPlanFragmentResult result = new PExecPlanFragmentResult();
+                            StatusPB pStatus = new StatusPB();
+                            pStatus.statusCode = TStatusCode.GLOBAL_DICT_NOT_MATCH.getValue();
+                            pStatus.errorMsgs = Collections.singletonList("test GLOBAL_DICT_NOT_MATCH error message");
+                            result.status = pStatus;
+                            return result;
+                        });
+                    }
+                };
+            }
+
+            return new MockPBackendService();
+        });
+        String sql =
+                "select count(1) from lineitem UNION ALL select count(1) from lineitem UNION ALL select count(1) from lineitem";
+        DefaultCoordinator scheduler = getScheduler(sql);
+        Assert.assertThrows(GlobalDictNotMatchException.class, scheduler::exec);
+    }
+
     @Test
     public void testDeployFutureReturnErrorStatus() throws Exception {
         final int successDeployedFragmentCount = 3;
diff --git a/test/sql/test_low_cardinality/R/test_low_cardinality_parquet_cancel b/test/sql/test_low_cardinality/R/test_low_cardinality_parquet_cancel
diff --git a/test/sql/test_low_cardinality/T/test_low_cardinality_parquet_cancel b/test/sql/test_low_cardinality/T/test_low_cardinality_parquet_cancel

Original file line number	Diff line number	Diff line change
`@@ -244,9 +244,12 @@ void FragmentContext::set_stream_load_contexts(const std::vector<StreamLoadConte`
`244`	`244`	`}`
`245`	`245`
`246`	`246`	`// Note: this function should be thread safe`
`247`		`-void FragmentContext::cancel(const Status& status) {`
	`247`	`+void FragmentContext::cancel(const Status& status, bool cancelled_by_fe) {`
`248`	`248`	`if (!status.ok() && _runtime_state != nullptr && _runtime_state->query_ctx() != nullptr) {`
`249`	`249`	`_runtime_state->query_ctx()->release_workgroup_token_once();`
	`250`	`+ if (cancelled_by_fe) {`
	`251`	`+ _runtime_state->query_ctx()->set_cancelled_by_fe();`
	`252`	`+ }`
`250`	`253`	`}`
`251`	`254`
`252`	`255`	`_runtime_state->set_is_cancelled(true);`
Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ class FragmentContext {`
`90`	`90`	`return status == nullptr ? Status::OK() : *status;`
`91`	`91`	`}`
`92`	`92`
`93`		`- void cancel(const Status& status);`
	`93`	`+ void cancel(const Status& status, bool cancelled_by_fe = false);`
`94`	`94`
`95`	`95`	`void finish() { cancel(Status::OK()); }`
`96`	`96`
Original file line number	Diff line number	Diff line change
`@@ -564,7 +564,7 @@ void PInternalServiceImplBase<T>::_cancel_plan_fragment(google::protobuf::RpcCon`
`564`	`564`	`"FragmentContext already destroyed: query_id=$0, fragment_instance_id=$1", print_id(query_id),`
`565`	`565`	`print_id(tid));`
`566`	`566`	`} else {`
`567`		`- fragment_ctx->cancel(Status::Cancelled(reason_string));`
	`567`	`+ fragment_ctx->cancel(Status::Cancelled(reason_string), true);`
`568`	`568`	`}`
`569`	`569`	`}`
`570`	`570`	`} else {`