feat(sharding): support sum function (#388)

* fix(sharding): fix min max error

Signed-off-by: wangbo <wangbo@sphere-ex.com>

* feat(sharding): support sum

Signed-off-by: wangbo <wangbo@sphere-ex.com>

* feat(sharding): support sum

Signed-off-by: wangbo <wangbo@sphere-ex.com>

* chore(sharding): delete space

Signed-off-by: wangbo <wangbo@sphere-ex.com>

Signed-off-by: wangbo <wangbo@sphere-ex.com>
Co-authored-by: wangbo <wangbo@sphere-ex.com>

Author: wbtlb

pisa-proxy/protocol/mysql/src/row.rs

@@ -28,7 +28,7 @@ pub trait RowData<T: AsRef<[u8]>> {
     fn get_row_data_with_name(&mut self, name: &str) -> value::Result<RowPartData>;
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum RowDataTyp<T: AsRef<[u8]>> {
     Text(RowDataText<T>),
     Binary(RowDataBinary<T>),
@@ -239,7 +239,6 @@ impl<T: AsRef<[u8]>> RowData<T> for RowDataBinary<T> {
 
                 // Need to add packet header and null_map to returnd data
                 let raw_data = &self.buf.as_ref()[start_pos + pos as usize..(start_pos + pos as usize + length as usize)];
-                println!("eeeeeeeeeeeee {:?}", &raw_data[..]);
                 return Ok(Some(
                     RowPartData { 
                         data: raw_data.into(), 

pisa-proxy/proxy/strategy/src/sharding_rewrite/meta.rs

@@ -39,6 +39,7 @@ pub enum FieldWrapFunc {
     Min,
     Max,
     Count,
+    Sum,
     None,
 }
 
@@ -48,6 +49,7 @@ impl AsRef<str> for FieldWrapFunc {
             Self::Max => "max",
             Self::Min => "min",
             Self::Count => "count",
+            Self::Sum => "sum",
             Self::None => "none",
         }
     }
@@ -378,6 +380,14 @@ impl Transformer for RewriteMetaData {
                                     );
                                 }
 
+                                AggFuncName::Sum => {
+                                    self.state = ScanState::FieldWrapFunc(
+                                        item.span,
+                                        FieldWrapFunc::Sum,
+                                        item.alias_name.clone(),
+                                    );
+                                }
+
                                 _ => {}
                             }
                             return false;

pisa-proxy/proxy/strategy/src/sharding_rewrite/mod.rs

@@ -1378,7 +1378,7 @@ impl ShardingRewrite {
                 .iter()
                 .filter_map(|f| {
                     if let FieldMeta::Ident(meta) = f {
-                        let is_match = matches!(meta.wrap_func,  FieldWrapFunc::Max| FieldWrapFunc::Min |FieldWrapFunc::Count);
+                        let is_match = matches!(meta.wrap_func,  FieldWrapFunc::Max| FieldWrapFunc::Min |FieldWrapFunc::Count| FieldWrapFunc::Sum);
                         return is_match.then(|| meta.clone())
                     } else {
                         None

pisa-proxy/runtime/mysql/src/server/executor.rs

@@ -215,7 +215,8 @@ where
                 .map_err(ErrorKind::from)?;
 
             let ro = &req.rewrite_outputs;
-            Self::handle_min_max(ro, &mut chunk, row_data.clone(), is_binary)?;
+
+            Self::handle_agg(ro, &mut chunk, row_data.clone(), is_binary)?;
 
             let avg_change = get_avg_change(&ro.results[0].changes);
 
@@ -260,7 +261,6 @@ where
 
                 let count: u64 = count_data.par_iter().sum();
                 let sum: u64 = sum_data.par_iter().sum();
-
                 chunk.par_iter_mut().for_each(|x| {
                     let mut row_data = row_data.clone();
                     row_data.with_buf(&x[4..]);
@@ -386,7 +386,7 @@ where
         Ok(())
     }
 
-    fn handle_min_max(
+    fn handle_agg(
         ro: &ShardingRewriteOutput,
         chunk: &mut [BytesMut],
         row_data: RowDataTyp<&[u8]>,
@@ -410,6 +410,54 @@ where
                         a.cmp(&b)
                     });
                 }
+                FieldWrapFunc::Sum => {
+                    let sum_data: Vec<_> = chunk
+                    .par_iter()
+                    .map(|x| -> Result<u64, Error> {
+                        let mut row_data = row_data.clone();
+                        row_data.with_buf(&x[4..]);
+                        let sum = if is_binary {
+                            let sum = decode_with_name::<&[u8], String>(
+                                &mut row_data,
+                                &agg.name,
+                                is_binary,
+                            ).unwrap();
+                            if let Some(sum) = sum {
+                                sum.parse::<u64>().map_err(|e| ErrorKind::Runtime(e.into()))?
+                            } else {
+                                0
+                            }
+                        } else {
+                            decode_with_name::<&[u8], u64>(&mut row_data, &agg.name, is_binary)
+                                .map_err(|e| ErrorKind::Runtime(e))?
+                                .unwrap_or_else(|| 0)
+                        };
+
+                        Ok(sum)
+                    })
+                    .collect::<Result<Vec<_>, _>>().unwrap();
+               
+                    let sum: u64 = sum_data.par_iter().sum();
+                    chunk.par_iter_mut().for_each(|x| {
+                        let mut row_data = row_data.clone();
+                        row_data.with_buf(&x[4..]);
+                        let sum_data =
+                            row_data.get_row_data_with_name(&agg.name).unwrap().unwrap();
+                        let part_data = RowPartData {
+                            data: vec![].into(),
+                            start_idx: sum_data.start_idx,
+                            part_encode_length: sum_data.part_encode_length,
+                            part_data_length: sum_data.part_data_length,
+                        };
+    
+                        let sum = format!("{:.4}", sum as u64);
+    
+                        row_data_cut_merge(x, &part_data, |data: &mut BytesMut| {
+                            data.put_lenc_int(sum.len() as u64, false);
+                            data.extend_from_slice(sum.as_bytes());
+                        });
+                    });
+                }
                 _ => {}
             }
 
@@ -709,5 +757,6 @@ fn get_min_max_value<'a>(
 
     row_data.with_buf(&b[4..]);
     let b = decode_with_name::<&[u8], u64>(row_data, name, is_binary).unwrap().unwrap();
+
     (a, b)
 }