copr: Vectorize uncompressed_length. (tikv#6109)

Signed-off-by: Zhu Zihao <all_but_last@163.com>
0x5457 · Dec 2, 2019 · b23731b · b23731b
1 parent 15f9af9
commit b23731b
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/components/tidb_query/src/rpn_expr/impl_encryption.rs b/components/tidb_query/src/rpn_expr/impl_encryption.rs
@@ -3,6 +3,8 @@
 use openssl::hash::{self, MessageDigest};
 use tidb_query_codegen::rpn_fn;
 
+use super::super::expr::{Error, EvalContext};
+
 use crate::codec::data_type::*;
 use crate::Result;
 
@@ -31,6 +33,22 @@ fn hex_digest(hashtype: MessageDigest, input: &[u8]) -> Result<Bytes> {
         .map_err(|e| box_err!("OpenSSL error: {:?}", e))
 }
 
+#[rpn_fn(capture = [ctx])]
+#[inline]
+pub fn uncompressed_length(ctx: &mut EvalContext, arg: &Option<Bytes>) -> Result<Option<Int>> {
+    use byteorder::{ByteOrder, LittleEndian};
+    Ok(arg.as_ref().map(|s| {
+        if s.is_empty() {
+            0
+        } else if s.len() <= 4 {
+            ctx.warnings.append_warning(Error::zlib_data_corrupted());
+            0
+        } else {
+            Int::from(LittleEndian::read_u32(&s[0..4]))
+        }
+    }))
+}
+
 #[cfg(test)]
 mod tests {
     use tipb::ScalarFuncSig;
@@ -53,6 +71,8 @@ mod tests {
         );
     }
 
+    use hex;
+
     #[test]
     fn test_md5() {
         let test_cases = vec![
@@ -117,4 +137,37 @@ mod tests {
         }
         test_unary_func_ok_none::<Bytes, Bytes>(ScalarFuncSig::Sha1);
     }
+
+    #[test]
+    fn test_uncompressed_length() {
+        let cases = vec![
+            (Some(""), Some(0)),
+            (
+                Some("0B000000789CCB48CDC9C95728CF2FCA4901001A0B045D"),
+                Some(11),
+            ),
+            (
+                Some("0C000000789CCB48CDC9C95728CF2F32303402001D8004202E"),
+                Some(12),
+            ),
+            (Some("020000000000"), Some(2)),
+            (Some("0000000001"), Some(0)),
+            (
+                Some("02000000789CCB48CDC9C95728CF2FCA4901001A0B045D"),
+                Some(2),
+            ),
+            (Some("010203"), Some(0)),
+            (Some("01020304"), Some(0)),
+            (None, None),
+        ];
+
+        for (s, exp) in cases {
+            let s = s.map(|inner| hex::decode(inner.as_bytes().to_vec()).unwrap());
+            let output = RpnFnScalarEvaluator::new()
+                .push_param(s)
+                .evaluate(ScalarFuncSig::UncompressedLength)
+                .unwrap();
+            assert_eq!(output, exp);
+        }
+    }
 }
diff --git a/components/tidb_query/src/rpn_expr/mod.rs b/components/tidb_query/src/rpn_expr/mod.rs
@@ -253,6 +253,7 @@ fn map_expr_node_to_rpn_func(expr: &Expr) -> Result<RpnFnMeta> {
         ScalarFuncSig::CaseWhenJson => case_when_fn_meta::<Json>(),
         ScalarFuncSig::Md5 => md5_fn_meta(),
         ScalarFuncSig::Sha1 => sha1_fn_meta(),
+        ScalarFuncSig::UncompressedLength => uncompressed_length_fn_meta(),
         ScalarFuncSig::DateFormatSig => date_format_fn_meta(),
         ScalarFuncSig::DayOfYear => day_of_year_fn_meta(),
         ScalarFuncSig::WeekDay => week_day_fn_meta(),