From 74809fad85a2a30e36c08e42c393c4785cb40c4e Mon Sep 17 00:00:00 2001
From: Drumil Patel <42701709+weastel@users.noreply.github.com>
Date: Tue, 2 Feb 2021 11:57:45 +0530
Subject: [PATCH] copr: Vectorize AddDatetimeAndString (#9610)

Signed-off-by: Drumil Patel <drumilpatel720@gmail.com>

<!--
Thank you for contributing to TiKV!

If you haven't already, please read TiKV's [CONTRIBUTING](https://github.com/tikv/tikv/blob/master/CONTRIBUTING.md) document.

If you're unsure about anything, just ask; somebody should be along to answer within a day or two.

PR Title Format:
1. module [, module2, module3]: what's changed
2. *: what's changed

If you want to open the **Challenge Program** pull request, please use the following template:
https://raw.githubusercontent.com/tikv/.github/master/.github/PULL_REQUEST_TEMPLATE/challenge-program.md
You can use it with query parameters: https://github.com/tikv/tikv/compare/master...${you branch}?template=challenge-program.md
-->

### What problem does this PR solve?

Issue Number: Partially fixes #9016 <!-- REMOVE this line if no issue to close -->

Problem Summary: Added Vectorize version for AddDatetimeAndString

### What is changed and how it works?

What's Changed: Implmented Vectorize vesion for AddDatetimeAndString

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

- Unit test

### Release note <!-- bugfixes or new feature need a release note -->
- Partially fixes #9016
---
 components/tidb_query_expr/src/impl_time.rs | 72 +++++++++++++++++++++
 components/tidb_query_expr/src/lib.rs       |  1 +
 2 files changed, 73 insertions(+)

diff --git a/components/tidb_query_expr/src/impl_time.rs b/components/tidb_query_expr/src/impl_time.rs
index bf0613b6da5..2c378c18eb1 100644
--- a/components/tidb_query_expr/src/impl_time.rs
+++ b/components/tidb_query_expr/src/impl_time.rs
@@ -296,6 +296,33 @@ pub fn add_datetime_and_duration(
     Ok(Some(res))
 }
 
+#[rpn_fn(capture=[ctx])]
+#[inline]
+pub fn add_datetime_and_string(
+    ctx: &mut EvalContext,
+    arg0: &DateTime,
+    arg1: BytesRef,
+) -> Result<Option<DateTime>> {
+    let arg1 = std::str::from_utf8(&arg1).map_err(Error::Encoding)?;
+    let arg1 = match Duration::parse(ctx, arg1, MAX_FSP) {
+        Ok(arg) => arg,
+        Err(_) => return Ok(None),
+    };
+
+    let res = match arg0.checked_add(ctx, arg1) {
+        Some(res) => res,
+        None => {
+            return ctx
+                .handle_invalid_time_error(Error::overflow(
+                    "DATETIME",
+                    format!("({} + {})", arg0, arg1),
+                ))
+                .map(|_| Ok(None))?;
+        }
+    };
+    Ok(Some(res))
+}
+
 #[rpn_fn(capture=[ctx])]
 #[inline]
 pub fn sub_duration_and_duration(
@@ -1312,6 +1339,51 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_add_datetime_and_string() {
+        let mut ctx = EvalContext::default();
+        let cases = vec![
+            // null cases
+            (None, None, None),
+            (None, Some("11:30:45.123456"), None),
+            (Some("2019-01-01 01:00:00"), None, None),
+            // normal cases
+            (
+                Some("2018-01-01"),
+                Some("11:30:45.123456"),
+                Some("2018-01-01 11:30:45.123456"),
+            ),
+            (
+                Some("2018-02-28 23:00:00"),
+                Some("01:30:30.123456"),
+                Some("2018-03-01 00:30:30.123456"),
+            ),
+            (
+                Some("2016-02-28 23:00:00"),
+                Some("01:30:30"),
+                Some("2016-02-29 00:30:30"),
+            ),
+            (
+                Some("2018-12-31 23:00:00"),
+                Some("01:30:30"),
+                Some("2019-01-01 00:30:30"),
+            ),
+        ];
+        for (arg0, arg1, exp) in cases {
+            let exp = exp.map(|exp| Time::parse_datetime(&mut ctx, exp, MAX_FSP, true).unwrap());
+            let arg0 =
+                arg0.map(|arg0| Time::parse_datetime(&mut ctx, arg0, MAX_FSP, true).unwrap());
+            let arg1 = arg1.map(|str| str.as_bytes().to_vec());
+
+            let output = RpnFnScalarEvaluator::new()
+                .push_param(arg0)
+                .push_param(arg1)
+                .evaluate(ScalarFuncSig::AddDatetimeAndString)
+                .unwrap();
+            assert_eq!(output, exp);
+        }
+    }
+
     #[test]
     fn test_sub_datetime_and_duration() {
         let mut ctx = EvalContext::default();
diff --git a/components/tidb_query_expr/src/lib.rs b/components/tidb_query_expr/src/lib.rs
index c55da2c79c5..14a98ae368d 100644
--- a/components/tidb_query_expr/src/lib.rs
+++ b/components/tidb_query_expr/src/lib.rs
@@ -614,6 +614,7 @@ fn map_expr_node_to_rpn_func(expr: &Expr) -> Result<RpnFnMeta> {
         ScalarFuncSig::ToSeconds => to_seconds_fn_meta(),
         ScalarFuncSig::DateDiff => date_diff_fn_meta(),
         ScalarFuncSig::AddDatetimeAndDuration => add_datetime_and_duration_fn_meta(),
+        ScalarFuncSig::AddDatetimeAndString => add_datetime_and_string_fn_meta(),
         ScalarFuncSig::SubDatetimeAndDuration => sub_datetime_and_duration_fn_meta(),
         ScalarFuncSig::FromDays => from_days_fn_meta(),
         ScalarFuncSig::Year => year_fn_meta(),