fix(processor): try fix test failure

Author: zhang2014

src/query/service/src/api/rpc/exchange/exchange_manager.rs

@@ -668,7 +668,8 @@ impl FragmentCoordinator {
 
             match &self.payload {
                 FragmentPayload::PlanV2(plan) => {
-                    let pipeline_builder = PipelineBuilderV2::create(ctx);
+                    let new_ctx = QueryContext::create_from(ctx);
+                    let pipeline_builder = PipelineBuilderV2::create(new_ctx);
                     self.pipeline_build_res = Some(pipeline_builder.finalize(plan)?);
                 }
             };

src/query/service/src/sessions/query_ctx.rs

@@ -69,7 +69,6 @@ pub struct QueryContext {
     version: String,
     partition_queue: Arc<RwLock<VecDeque<PartInfoPtr>>>,
     shared: Arc<QueryContextShared>,
-    precommit_blocks: Arc<RwLock<Vec<DataBlock>>>,
     fragment_id: Arc<AtomicUsize>,
 }
 
@@ -85,7 +84,6 @@ impl QueryContext {
             partition_queue: Arc::new(RwLock::new(VecDeque::new())),
             version: format!("DatabendQuery {}", *crate::version::DATABEND_COMMIT_VERSION),
             shared,
-            precommit_blocks: Arc::new(RwLock::new(Vec::new())),
             fragment_id: Arc::new(AtomicUsize::new(0)),
         })
     }
@@ -334,15 +332,10 @@ impl TableContext for QueryContext {
         self.shared.dal_ctx.as_ref()
     }
     fn push_precommit_block(&self, block: DataBlock) {
-        let mut blocks = self.precommit_blocks.write();
-        blocks.push(block);
+        self.shared.push_precommit_block(block)
     }
     fn consume_precommit_blocks(&self) -> Vec<DataBlock> {
-        let mut blocks = self.precommit_blocks.write();
-
-        let mut swaped_precommit_blocks = vec![];
-        std::mem::swap(&mut *blocks, &mut swaped_precommit_blocks);
-        swaped_precommit_blocks
+        self.shared.consume_precommit_blocks()
     }
     fn try_get_function_context(&self) -> Result<FunctionContext> {
         let tz = self.get_settings().get_timezone()?;

src/query/service/src/sessions/query_ctx_shared.rs

@@ -22,6 +22,7 @@ use std::sync::Weak;
 use common_base::base::Progress;
 use common_base::base::Runtime;
 use common_contexts::DalContext;
+use common_datablocks::DataBlock;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_types::UserInfo;
@@ -79,6 +80,7 @@ pub struct QueryContextShared {
     pub(in crate::sessions) catalog_manager: Arc<CatalogManager>,
     pub(in crate::sessions) storage_operator: StorageOperator,
     pub(in crate::sessions) executor: Arc<RwLock<Weak<PipelineExecutor>>>,
+    pub(in crate::sessions) precommit_blocks: Arc<RwLock<Vec<DataBlock>>>,
 }
 
 impl QueryContextShared {
@@ -108,6 +110,7 @@ impl QueryContextShared {
             auth_manager: AuthMgr::create(config).await?,
             affect: Arc::new(Mutex::new(None)),
             executor: Arc::new(RwLock::new(Weak::new())),
+            precommit_blocks: Arc::new(RwLock::new(vec![])),
         }))
     }
 
@@ -297,4 +300,17 @@ impl QueryContextShared {
         let mut executor = self.executor.write();
         *executor = weak_ptr;
     }
+
+    pub fn push_precommit_block(&self, block: DataBlock) {
+        let mut blocks = self.precommit_blocks.write();
+        blocks.push(block);
+    }
+
+    pub fn consume_precommit_blocks(&self) -> Vec<DataBlock> {
+        let mut blocks = self.precommit_blocks.write();
+
+        let mut swaped_precommit_blocks = vec![];
+        std::mem::swap(&mut *blocks, &mut swaped_precommit_blocks);
+        swaped_precommit_blocks
+    }
 }