Merge pull request #2042 from GrapeBaBa/tomonth

[Feature]Add toMonth date function
databendlabs · Oct 1, 2021 · 95e43bf · 95e43bf
2 parents 8aa6b98 + 7d707d7
commit 95e43bf
Show file tree

Hide file tree

Showing 6 changed files with 182 additions and 2 deletions.
diff --git a/common/functions/src/scalars/dates/date.rs b/common/functions/src/scalars/dates/date.rs
@@ -19,6 +19,7 @@ use super::interval_function::MonthsArithmeticFunction;
 use super::interval_function::SecondsArithmeticFunction;
 use super::now::NowFunction;
 use super::RoundFunction;
+use super::ToMonthFunction;
 use super::ToStartOfISOYearFunction;
 use super::ToStartOfMonthFunction;
 use super::ToStartOfQuarterFunction;
@@ -59,7 +60,7 @@ impl DateFunction {
         );
         map.insert("toStartOfWeek".into(), ToStartOfWeekFunction::try_create);
         map.insert("toStartOfMonth".into(), ToStartOfMonthFunction::try_create);
-
+        map.insert("toMonth".into(), ToMonthFunction::try_create);
         // rounders
         {
             map.insert("toStartOfSecond".into(), |display_name| {

diff --git a/common/functions/src/scalars/dates/date_function_test.rs b/common/functions/src/scalars/dates/date_function_test.rs
@@ -14,10 +14,12 @@
 
 use common_arrow::arrow::array::UInt32Array;
 use common_arrow::arrow::array::UInt64Array;
+use common_arrow::arrow::array::UInt8Array;
 use common_datablocks::DataBlock;
 use common_datavalues::prelude::*;
 use common_exception::Result;
 
+use crate::scalars::dates::number_function::ToMonthFunction;
 use crate::scalars::ToYYYYMMDDFunction;
 use crate::scalars::ToYYYYMMDDhhmmssFunction;
 use crate::scalars::ToYYYYMMFunction;
@@ -538,3 +540,148 @@ fn test_toyyyymmhhmmss_constant_function() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_tomonth_constant_function() -> Result<()> {
+    // date16
+    let schema = DataSchemaRefExt::create(vec![DataField::new("a", DataType::Date16, false)]);
+    let block = DataBlock::create(schema.clone(), vec![DataColumn::Constant(
+        DataValue::UInt16(Some(0u16)),
+        5,
+    )]);
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 5);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([1; 5]);
+
+        assert_eq!(actual, &expected);
+    }
+
+    // date32
+    let schema = DataSchemaRefExt::create(vec![DataField::new("a", DataType::Date32, false)]);
+    let block = DataBlock::create(schema.clone(), vec![DataColumn::Constant(
+        DataValue::Int32(Some(0i32)),
+        10,
+    )]);
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 10);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([1; 10]);
+
+        assert_eq!(actual, &expected);
+    }
+
+    // datetime
+    // 2021-10-01 17:50:17 --- 1633081817
+    let schema =
+        DataSchemaRefExt::create(vec![DataField::new("a", DataType::DateTime32(None), false)]);
+    let block = DataBlock::create(schema.clone(), vec![DataColumn::Constant(
+        DataValue::UInt32(Some(1633081817u32)),
+        15,
+    )]);
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 15);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([10; 15]);
+
+        assert_eq!(actual, &expected);
+    }
+
+    Ok(())
+}
+
+#[test]
+fn test_tomonth_function() -> Result<()> {
+    // date16
+    let schema = DataSchemaRefExt::create(vec![DataField::new("a", DataType::Date16, false)]);
+    let block = DataBlock::create_by_array(schema.clone(), vec![Series::new(vec![0u16])]);
+
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 1);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([1; 1]);
+        assert_eq!(actual, &expected);
+    }
+
+    // date32
+    let schema = DataSchemaRefExt::create(vec![DataField::new("a", DataType::Date32, false)]);
+    let block = DataBlock::create_by_array(schema.clone(), vec![Series::new(vec![0i32])]);
+
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 1);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([1; 1]);
+
+        assert_eq!(actual, &expected);
+    }
+
+    // dateTime
+    // 2021-10-01 17:50:17 --- 1633081817
+    let schema =
+        DataSchemaRefExt::create(vec![DataField::new("a", DataType::DateTime32(None), false)]);
+    let block = DataBlock::create_by_array(schema.clone(), vec![Series::new(vec![1633081817u32])]);
+
+    {
+        let col = ToMonthFunction::try_create("a")?;
+        let columns = vec![DataColumnWithField::new(
+            block.try_column_by_name("a")?.clone(),
+            schema.field_with_name("a")?.clone(),
+        )];
+        let result = col.eval(&columns, block.num_rows())?;
+        assert_eq!(result.len(), 1);
+        assert_eq!(result.data_type(), DataType::UInt8);
+
+        let actual_ref = result.get_array_ref().unwrap();
+        let actual = actual_ref.as_any().downcast_ref::<UInt8Array>().unwrap();
+        let expected = UInt8Array::from_slice([10; 1]);
+
+        assert_eq!(actual, &expected);
+    }
+
+    Ok(())
+}
diff --git a/common/functions/src/scalars/dates/mod.rs b/common/functions/src/scalars/dates/mod.rs
@@ -33,6 +33,7 @@ pub use interval_function::IntervalArithmeticFunction;
 pub use interval_function::IntervalFunctionFactory;
 pub use interval_function::MonthsArithmeticFunction;
 pub use interval_function::SecondsArithmeticFunction;
+pub use number_function::ToMonthFunction;
 pub use number_function::ToStartOfISOYearFunction;
 pub use number_function::ToStartOfMonthFunction;
 pub use number_function::ToStartOfQuarterFunction;

diff --git a/common/functions/src/scalars/dates/number_function.rs b/common/functions/src/scalars/dates/number_function.rs
@@ -167,6 +167,22 @@ impl NumberResultFunction<u16> for ToStartOfMonth {
     }
 }
 
+#[derive(Clone)]
+pub struct ToMonth;
+
+impl NumberResultFunction<u8> for ToMonth {
+    fn return_type() -> Result<DataType> {
+        Ok(DataType::UInt8)
+    }
+    fn to_number(value: DateTime<Utc>) -> u8 {
+        value.month() as u8
+    }
+
+    fn to_constant_value(value: DateTime<Utc>) -> DataValue {
+        DataValue::UInt8(Some(Self::to_number(value)))
+    }
+}
+
 impl<T, R> NumberFunction<T, R>
 where
     T: NumberResultFunction<R> + Clone + Sync + Send + 'static,
@@ -287,3 +303,5 @@ pub type ToStartOfISOYearFunction = NumberFunction<ToStartOfISOYear, u16>;
 pub type ToStartOfYearFunction = NumberFunction<ToStartOfYear, u16>;
 pub type ToStartOfQuarterFunction = NumberFunction<ToStartOfQuarter, u16>;
 pub type ToStartOfMonthFunction = NumberFunction<ToStartOfMonth, u16>;
+
+pub type ToMonthFunction = NumberFunction<ToMonth, u8>;
diff --git a/tests/suites/0_stateless/02_0012_function_datetimes.result b/tests/suites/0_stateless/02_0012_function_datetimes.result
@@ -106,3 +106,9 @@
 ===addSeconds===
 2020-02-29 10:01:01
 ===addSeconds===
+===toMonth===
+10
+10
+1
+1
+===toMonth===
diff --git a/tests/suites/0_stateless/02_0012_function_datetimes.sql b/tests/suites/0_stateless/02_0012_function_datetimes.sql
@@ -123,4 +123,11 @@ select '===subtractMinutes===';
 
 select '===addSeconds===';
 select addSeconds(toDateTime(1582970400), cast(61, INT32)); -- 2020-2-29T10:00:00 + 61 seconds
-select '===addSeconds===';
+select '===addSeconds===';
+
+select '===toMonth===';
+select toMonth(toDateTime(1633081817));
+select toMonth(toDate(18901));
+select toMonth(toDateTime(1633081817))  =  10;
+select toMonth(toDate(18901))  =  10;
+select '===toMonth===';