Merge pull request #9251 from zhyass/feature_cluster

feat(query): insert/delete support display rows affected on mysql client

Author: mergify[bot]

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -469,4 +469,9 @@ impl ProgressReporter for ContextProgressReporter {
             convert_byte_size((progress.bytes as f64) / (seconds)),
         )
     }
+
+    fn affected_rows(&self) -> u64 {
+        let progress = self.context.get_scan_progress_value();
+        progress.rows as u64
+    }
 }

src/query/service/src/servers/mysql/writers/query_result_writer.rs

@@ -38,6 +38,7 @@ use tracing::error;
 /// "Read x rows, y MiB in z sec., A million rows/sec., B MiB/sec."
 pub trait ProgressReporter {
     fn progress_info(&self) -> String;
+    fn affected_rows(&self) -> u64;
 }
 
 pub struct QueryResult {
@@ -122,7 +123,16 @@ impl<'a, W: AsyncWrite + Send + Unpin> DFQueryResultWriter<'a, W> {
                 }
             }
 
-            dataset_writer.completed(OkResponse::default()).await?;
+            let affected_rows = query_result
+                .extra_info
+                .map(|r| r.affected_rows())
+                .unwrap_or_default();
+            dataset_writer
+                .completed(OkResponse {
+                    affected_rows,
+                    ..Default::default()
+                })
+                .await?;
             return Ok(());
         }
 

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

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use common_base::base::ProgressValues;
 use common_catalog::plan::Expression;
 use common_catalog::plan::Partitions;
 use common_catalog::plan::PartitionsShuffleKind;
@@ -74,8 +75,14 @@ impl FuseTable {
             return Ok(());
         }
 
+        let scan_progress = ctx.get_scan_progress();
         // check if unconditional deletion
         if filter.is_none() {
+            let progress_values = ProgressValues {
+                rows: snapshot.summary.row_count as usize,
+                bytes: snapshot.summary.uncompressed_byte_size as usize,
+            };
+            scan_progress.incr(&progress_values);
             // deleting the whole table... just a truncate
             let purge = false;
             return self.do_truncate(ctx.clone(), purge).await;
@@ -91,7 +98,19 @@ impl FuseTable {
             // if the `filter_expr` is of "constant" nullary :
             //   for the whole block, whether all of the rows should be kept or dropped,
             //   we can just return from here, without accessing the block data
-            return self.delete_eval_const(ctx.clone(), &filter_expr).await;
+            if self.try_eval_const(&filter_expr)? {
+                let progress_values = ProgressValues {
+                    rows: snapshot.summary.row_count as usize,
+                    bytes: snapshot.summary.uncompressed_byte_size as usize,
+                };
+                scan_progress.incr(&progress_values);
+
+                // deleting the whole table... just a truncate
+                let purge = false;
+                return self.do_truncate(ctx.clone(), purge).await;
+            }
+            // do nothing.
+            return Ok(());
         }
 
         self.try_add_deletion_source(ctx.clone(), &filter_expr, col_indices, &snapshot, pipeline)
@@ -105,32 +124,21 @@ impl FuseTable {
         Ok(())
     }
 
-    async fn delete_eval_const(
-        &self,
-        ctx: Arc<dyn TableContext>,
-        filter: &Expression,
-    ) -> Result<()> {
+    fn try_eval_const(&self, filter: &Expression) -> Result<bool> {
         let func_ctx = FunctionContext::default();
 
         let dummy_field = DataField::new("dummy", NullType::new_impl());
-        let dummy_schema = Arc::new(DataSchema::new(vec![dummy_field.clone()]));
+        let dummy_schema = Arc::new(DataSchema::new(vec![dummy_field]));
         let dummy_column = DataValue::Null.as_const_column(&NullType::new_impl(), 1)?;
         let dummy_data_block = DataBlock::create(dummy_schema.clone(), vec![dummy_column]);
 
         let eval_node = Arc::new(Evaluator::eval_expression(filter, dummy_schema.as_ref())?);
         let filter_result = eval_node.eval(&func_ctx, &dummy_data_block)?.vector;
         debug_assert!(filter_result.len() == 1);
 
-        let filter_result = DataBlock::cast_to_nonull_boolean(&filter_result)?
+        DataBlock::cast_to_nonull_boolean(&filter_result)?
             .get(0)
-            .as_bool()?;
-        if filter_result {
-            // deleting the whole table... just a truncate
-            let purge = false;
-            return self.do_truncate(ctx.clone(), purge).await;
-        }
-        // do nothing.
-        Ok(())
+            .as_bool()
     }
 
     async fn try_add_deletion_source(

src/query/storages/fuse/fuse/src/operations/mutation/mutation_sink.rs

@@ -15,6 +15,8 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use common_base::base::Progress;
+use common_base::base::ProgressValues;
 use common_catalog::table::Table;
 use common_catalog::table::TableExt;
 use common_catalog::table_context::TableContext;
@@ -63,6 +65,7 @@ pub struct MutationSink {
     ctx: Arc<dyn TableContext>,
     dal: Operator,
     location_gen: TableMetaLocationGenerator,
+    scan_progress: Arc<Progress>,
 
     table: Arc<dyn Table>,
     base_snapshot: Arc<TableSnapshot>,
@@ -84,11 +87,13 @@ impl MutationSink {
         base_snapshot: Arc<TableSnapshot>,
         input: Arc<InputPort>,
     ) -> Result<ProcessorPtr> {
+        let scan_progress = ctx.get_scan_progress();
         Ok(ProcessorPtr::create(Box::new(MutationSink {
             state: State::None,
             ctx,
             dal: table.get_operator(),
             location_gen: table.meta_location_generator.clone(),
+            scan_progress,
             table: Arc::new(table.clone()),
             base_snapshot,
             merged_segments: vec![],
@@ -154,6 +159,23 @@ impl Processor for MutationSink {
         match std::mem::replace(&mut self.state, State::None) {
             State::ReadMeta(input_meta) => {
                 let meta = MutationMeta::from_meta(&input_meta)?;
+
+                let affect_rows = self
+                    .base_snapshot
+                    .summary
+                    .row_count
+                    .abs_diff(meta.summary.row_count);
+                let affect_bytes = self
+                    .base_snapshot
+                    .summary
+                    .uncompressed_byte_size
+                    .abs_diff(meta.summary.uncompressed_byte_size);
+                let progress_values = ProgressValues {
+                    rows: affect_rows as usize,
+                    bytes: affect_bytes as usize,
+                };
+                self.scan_progress.incr(&progress_values);
+
                 self.merged_segments = meta.segments.clone();
                 self.merged_statistics = meta.summary.clone();
                 self.abort_operation = meta.abort_operation.clone();