chore: remove spawn_blocking for merge into and replace into (#14664)

remove spawn_blocking for merge into and replace into

Author: JackTan25

src/query/storages/fuse/src/operations/merge_into/mutator/matched_mutator.rs

@@ -337,8 +337,10 @@ impl MatchedAggregator {
         let log_entries = futures::future::try_join_all(mutation_log_handlers)
             .await
             .map_err(|e| {
-                ErrorCode::Internal("unexpected, failed to join apply-deletion tasks.")
-                    .add_message_back(e.to_string())
+                ErrorCode::Internal(
+                    "unexpected, failed to join apply update and delete tasks for merge into.",
+                )
+                .add_message_back(e.to_string())
             })?;
         let mut mutation_logs = Vec::new();
         for maybe_log_entry in log_entries {
@@ -377,6 +379,7 @@ impl AggregationContext {
             &self.block_reader,
             block_meta,
             &self.read_settings,
+            self.ctx.get_id(),
         )
         .await?;
         let origin_num_rows = origin_data_block.num_rows();
@@ -407,8 +410,9 @@ impl AggregationContext {
         // and wait (asyncly, which will NOT block the executor thread)
         let block_builder = self.block_builder.clone();
         let origin_stats = block_meta.cluster_stats.clone();
+
         let serialized = GlobalIORuntime::instance()
-            .spawn_blocking(move || {
+            .spawn(self.ctx.get_id(), async move {
                 block_builder.build(res_block, |block, generator| {
                     let cluster_stats =
                         generator.gen_with_origin_stats(&block, origin_stats.clone())?;
@@ -419,7 +423,13 @@ impl AggregationContext {
                     Ok((cluster_stats, block))
                 })
             })
-            .await?;
+            .await
+            .map_err(|e| {
+                ErrorCode::Internal(
+                    "unexpected, failed to serialize block when apply update and delete to data block for merge into",
+                )
+                .add_message_back(e.to_string())
+            })??;
 
         // persistent data
         let new_block_meta = serialized.block_meta;

src/query/storages/fuse/src/operations/replace_into/mutator/merge_into_mutator.rs

@@ -89,6 +89,7 @@ struct AggregationContext {
     segment_reader: CompactSegmentInfoReader,
     block_builder: BlockBuilder,
     io_request_semaphore: Arc<Semaphore>,
+    query_id: String,
 }
 
 // Apply MergeIntoOperations to segments
@@ -166,7 +167,7 @@ impl MergeIntoOperationAggregator {
                 Some(reader)
             }
         };
-
+        let query_id = ctx.get_id();
         Ok(Self {
             ctx,
             deletion_accumulator,
@@ -184,6 +185,7 @@ impl MergeIntoOperationAggregator {
                 segment_reader,
                 block_builder,
                 io_request_semaphore,
+                query_id,
             }),
         })
     }
@@ -392,6 +394,7 @@ impl AggregationContext {
             &self.key_column_reader,
             block_meta,
             &self.read_settings,
+            self.query_id.clone(),
         )
         .await?;
 
@@ -507,14 +510,20 @@ impl AggregationContext {
         let block_builder = self.block_builder.clone();
         let origin_stats = block_meta.cluster_stats.clone();
         let serialized = GlobalIORuntime::instance()
-            .spawn_blocking(move || {
+            .spawn(self.query_id.clone(), async move {
                 block_builder.build(new_block, |block, generator| {
                     let cluster_stats =
                         generator.gen_with_origin_stats(&block, origin_stats.clone())?;
                     Ok((cluster_stats, block))
                 })
             })
-            .await?;
+            .await
+            .map_err(|e| {
+                ErrorCode::Internal(
+                    "unexpected, failed to join apply delete tasks for replace into.",
+                )
+                .add_message_back(e.to_string())
+            })??;
 
         // persistent data
         let new_block_meta = serialized.block_meta;
@@ -600,7 +609,7 @@ impl AggregationContext {
         let block_meta_ptr = block_meta.clone();
         let reader = reader.clone();
         GlobalIORuntime::instance()
-            .spawn_blocking(move || {
+            .spawn(self.query_id.clone(), async move {
                 let column_chunks = merged_io_read_result.columns_chunks()?;
                 reader.deserialize_chunks(
                     block_meta_ptr.location.0.as_str(),
@@ -612,6 +621,12 @@ impl AggregationContext {
                 )
             })
             .await
+            .map_err(|e| {
+                ErrorCode::Internal(
+                    "unexpected, failed to join aggregation context read block tasks for replace into.",
+                )
+                .add_message_back(e.to_string())
+            })?
     }
 
     // return true if the block is pruned, otherwise false

src/query/storages/fuse/src/operations/util.rs

@@ -22,6 +22,7 @@ use databend_common_arrow::parquet::metadata::ThriftFileMetaData;
 use databend_common_base::base::tokio::sync::OwnedSemaphorePermit;
 use databend_common_base::base::tokio::sync::Semaphore;
 use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_base::runtime::TrySpawn;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::ColumnId;
@@ -199,6 +200,7 @@ pub async fn read_block(
     reader: &BlockReader,
     block_meta: &BlockMeta,
     read_settings: &ReadSettings,
+    query_id: String,
 ) -> Result<DataBlock> {
     let merged_io_read_result = reader
         .read_columns_data_by_merge_io(
@@ -211,11 +213,11 @@ pub async fn read_block(
 
     // deserialize block data
     // cpu intensive task, send them to dedicated thread pool
-
     let block_meta_ptr = block_meta.clone();
     let reader = reader.clone();
+
     GlobalIORuntime::instance()
-        .spawn_blocking(move || {
+        .spawn(query_id, async move {
             let column_chunks = merged_io_read_result.columns_chunks()?;
             reader.deserialize_chunks(
                 block_meta_ptr.location.0.as_str(),
@@ -227,4 +229,8 @@ pub async fn read_block(
             )
         })
         .await
+        .map_err(|e| {
+            ErrorCode::Internal("unexpected, failed to read block for merge into")
+                .add_message_back(e.to_string())
+        })?
 }