!5849 修复查询中设置用户变量卡住的问题

Merge pull request !5849 from TinyBag/uservar

src/gausskernel/runtime/executor/nodeIndexscan.cpp

@@ -168,7 +168,7 @@ static TupleTableSlot* ExecIndexScan(PlanState* state)
     /*
      * If we have runtime keys and they've not already been set up, do it now.
      */
-    if (node->iss_NumRuntimeKeys != 0 && (!node->iss_RuntimeKeysReady || (u_sess->parser_cxt.has_set_uservar && DB_IS_CMPT(B_FORMAT)))) {
+    if (node->iss_NumRuntimeKeys != 0 && !node->iss_RuntimeKeysReady) {
         /*
          * set a flag for partitioned table, so we can deal with it specially
          * when we rescan the partitioned table
@@ -181,6 +181,10 @@ static TupleTableSlot* ExecIndexScan(PlanState* state)
         } else {
             ExecReScan((PlanState*)node);
         }
+    } else if (DB_IS_CMPT(B_FORMAT) && node->iss_NumRuntimeKeys != 0 && u_sess->parser_cxt.has_set_uservar) {
+        ExprContext* econtext = node->iss_RuntimeContext;
+        ResetExprContext(econtext);
+        ExecIndexEvalRuntimeKeys(econtext, node->iss_RuntimeKeys, node->iss_NumRuntimeKeys);
     }
 
     return ExecScan(&node->ss, (ExecScanAccessMtd)IndexNext, (ExecScanRecheckMtd)IndexRecheck);

src/gausskernel/runtime/executor/nodeSubplan.cpp

@@ -991,8 +991,10 @@ void ExecSetParamPlan(SubPlanState* node, ExprContext* econtext)
     MemoryContext oldcontext = MemoryContextSwitchTo(econtext->ecxt_per_query_memory);
 
     if (u_sess->parser_cxt.has_set_uservar && DB_IS_CMPT(B_FORMAT)) {
-        if (nodeTag(planstate) == T_SeqScanState) {
+        if (IsA(planstate, SeqScanState)) {
             scan_handler_tbl_restrpos(castNode(SeqScanState, planstate)->ss_currentScanDesc);
+        } else if (IsA(planstate, IndexScanState)) {
+            ExecReScan(planstate);
         }
     }
 

src/test/regress/input/set_user_defined_variables_test.source

@@ -774,6 +774,23 @@ END WHILE label_1;
 end;
 /
 SELECT TRIM(TRAILING ', ' FROM @sequence);
+
+-- index scan卡住问题
+\c test_set
+set enable_set_variable_b_format = 1;
+set enable_seqscan = false;
+set enable_bitmapscan = false;
+create table account_instance_stats(id int primary key, stat_id int);
+insert into account_instance_stats values(1,1),(2,1),(3,1),(4,1);
+set @pids = 1;
+
+SELECT stat_id, @pids := id FROM account_instance_stats where id >= @pids;
+
+drop table account_instance_stats;
+set enable_seqscan = default;
+set enable_bitmapscan = default;
+set enable_set_variable_b_format = default;
+
 \c regression
 drop database if exists test_set;
 

src/test/regress/output/set_user_defined_variables_test.source

@@ -1520,6 +1520,28 @@ SELECT TRIM(TRAILING ', ' FROM @sequence);
  1, 2, 3, 4, 5, 6, 7, 8, 9, 10
 (1 row)
 
+-- index scan卡住问题
+\c test_set
+set enable_set_variable_b_format = 1;
+set enable_seqscan = false;
+set enable_bitmapscan = false;
+create table account_instance_stats(id int primary key, stat_id int);
+NOTICE:  CREATE TABLE / PRIMARY KEY will create implicit index "account_instance_stats_pkey" for table "account_instance_stats"
+insert into account_instance_stats values(1,1),(2,1),(3,1),(4,1);
+set @pids = 1;
+SELECT stat_id, @pids := id FROM account_instance_stats where id >= @pids;
+ stat_id | ?column? 
+---------+----------
+       1 |        1
+       1 |        2
+       1 |        3
+       1 |        4
+(4 rows)
+
+drop table account_instance_stats;
+set enable_seqscan = default;
+set enable_bitmapscan = default;
+set enable_set_variable_b_format = default;
 \c regression
 drop database if exists test_set;
 \! @abs_bindir@/gs_guc reload -Z datanode -D @abs_srcdir@/tmp_check/datanode1 -c "enable_set_variable_b_format=off" >/dev/null 2>&1