diff --git a/src/common/function/src/helper.rs b/src/common/function/src/helper.rs new file mode 100644 index 000000000000..6f549d6619e3 --- /dev/null +++ b/src/common/function/src/helper.rs @@ -0,0 +1,29 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use common_query::prelude::{Signature, TypeSignature, Volatility}; +use datatypes::prelude::ConcreteDataType; + +/// Create a function signature with oneof signatures of interleaving two arguments. +pub fn one_of_sigs2(args1: Vec, args2: Vec) -> Signature { + let mut sigs = Vec::with_capacity(args1.len() * args2.len()); + + for arg1 in &args1 { + for arg2 in &args2 { + sigs.push(TypeSignature::Exact(vec![arg1.clone(), arg2.clone()])); + } + } + + Signature::one_of(sigs, Volatility::Immutable) +} diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index a971d5cceb0d..5d3ab6d42069 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -13,3 +13,5 @@ // limitations under the License. pub mod scalars; + +pub mod helper; diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index e6f7ed766c95..3aa19632e7df 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod aggregate; +mod date; pub mod expression; pub mod function; pub mod function_registry; diff --git a/src/common/function/src/scalars/date.rs b/src/common/function/src/scalars/date.rs new file mode 100644 index 000000000000..0e16019d527b --- /dev/null +++ b/src/common/function/src/scalars/date.rs @@ -0,0 +1,31 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +mod date_add; +mod date_sub; + +use date_add::DateAddFunction; +use date_sub::DateSubFunction; + +use crate::scalars::function_registry::FunctionRegistry; + +pub(crate) struct DateFunction; + +impl DateFunction { + pub fn register(registry: &FunctionRegistry) { + registry.register(Arc::new(DateAddFunction)); + registry.register(Arc::new(DateSubFunction)); + } +} diff --git a/src/common/function/src/scalars/date/date_add.rs b/src/common/function/src/scalars/date/date_add.rs new file mode 100644 index 000000000000..e299f7947297 --- /dev/null +++ b/src/common/function/src/scalars/date/date_add.rs @@ -0,0 +1,279 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt; + +use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu}; +use common_query::prelude::Signature; +use datatypes::data_type::DataType; +use datatypes::prelude::ConcreteDataType; +use datatypes::value::ValueRef; +use datatypes::vectors::VectorRef; +use snafu::ensure; + +use crate::helper; +use crate::scalars::function::{Function, FunctionContext}; + +/// A function adds an interval value to Timestamp, Date or DateTime, and return the result. +#[derive(Clone, Debug, Default)] +pub struct DateAddFunction; + +const NAME: &str = "date_add"; + +impl Function for DateAddFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + Ok(input_types[0].clone()) + } + + fn signature(&self) -> Signature { + helper::one_of_sigs2( + vec![ + ConcreteDataType::date_datatype(), + ConcreteDataType::datetime_datatype(), + ConcreteDataType::timestamp_second_datatype(), + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_microsecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ], + vec![ + ConcreteDataType::interval_month_day_nano_datatype(), + ConcreteDataType::interval_year_month_datatype(), + ConcreteDataType::interval_day_time_datatype(), + ], + ) + } + + fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result { + ensure!( + columns.len() == 2, + InvalidFuncArgsSnafu { + err_msg: format!( + "The length of the args is not correct, expect 2, have: {}", + columns.len() + ), + } + ); + + let left = &columns[0]; + let right = &columns[1]; + + let size = left.len(); + let left_datatype = columns[0].data_type(); + match left_datatype { + ConcreteDataType::Timestamp(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let ts = left.get(i).as_timestamp(); + let interval = right.get(i).as_interval(); + + let new_ts = match (ts, interval) { + (Some(ts), Some(interval)) => ts.add_interval(interval), + _ => ts, + }; + + result.push_value_ref(ValueRef::from(new_ts)); + } + + Ok(result.to_vector()) + } + ConcreteDataType::Date(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let date = left.get(i).as_date(); + let interval = right.get(i).as_interval(); + let new_date = match (date, interval) { + (Some(date), Some(interval)) => date.add_interval(interval), + _ => date, + }; + + result.push_value_ref(ValueRef::from(new_date)); + } + + Ok(result.to_vector()) + } + ConcreteDataType::DateTime(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let datetime = left.get(i).as_datetime(); + let interval = right.get(i).as_interval(); + let new_datetime = match (datetime, interval) { + (Some(datetime), Some(interval)) => datetime.add_interval(interval), + _ => datetime, + }; + + result.push_value_ref(ValueRef::from(new_datetime)); + } + + Ok(result.to_vector()) + } + _ => UnsupportedInputDataTypeSnafu { + function: NAME, + datatypes: columns.iter().map(|c| c.data_type()).collect::>(), + } + .fail(), + } + } +} + +impl fmt::Display for DateAddFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "DATE_ADD") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_query::prelude::{TypeSignature, Volatility}; + use datatypes::prelude::ConcreteDataType; + use datatypes::value::Value; + use datatypes::vectors::{ + DateTimeVector, DateVector, IntervalDayTimeVector, IntervalYearMonthVector, + TimestampSecondVector, + }; + + use super::{DateAddFunction, *}; + use crate::scalars::Function; + + #[test] + fn test_date_add_misc() { + let f = DateAddFunction; + assert_eq!("date_add", f.name()); + assert_eq!( + ConcreteDataType::timestamp_microsecond_datatype(), + f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()]) + .unwrap() + ); + assert_eq!( + ConcreteDataType::timestamp_second_datatype(), + f.return_type(&[ConcreteDataType::timestamp_second_datatype()]) + .unwrap() + ); + assert_eq!( + ConcreteDataType::date_datatype(), + f.return_type(&[ConcreteDataType::date_datatype()]).unwrap() + ); + assert_eq!( + ConcreteDataType::datetime_datatype(), + f.return_type(&[ConcreteDataType::datetime_datatype()]) + .unwrap() + ); + assert!(matches!(f.signature(), + Signature { + type_signature: TypeSignature::OneOf(sigs), + volatility: Volatility::Immutable + } if sigs.len() == 18)); + } + + #[test] + fn test_timestamp_date_add() { + let f = DateAddFunction; + + let times = vec![Some(123), None, Some(42), None]; + // Intervals in milliseconds + let intervals = vec![1000, 2000, 3000, 1000]; + let results = [Some(124), None, Some(45), None]; + + let time_vector = TimestampSecondVector::from(times.clone()); + let interval_vector = IntervalDayTimeVector::from_vec(intervals); + let args: Vec = vec![Arc::new(time_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in times.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::Timestamp(ts) => { + assert_eq!(ts.value(), result.unwrap()); + } + _ => unreachable!(), + } + } + } + + #[test] + fn test_date_date_add() { + let f = DateAddFunction; + + let dates = vec![Some(123), None, Some(42), None]; + // Intervals in months + let intervals = vec![1, 2, 3, 1]; + let results = [Some(154), None, Some(131), None]; + + let date_vector = DateVector::from(dates.clone()); + let interval_vector = IntervalYearMonthVector::from_vec(intervals); + let args: Vec = vec![Arc::new(date_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in dates.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::Date(date) => { + assert_eq!(date.val(), result.unwrap()); + } + _ => unreachable!(), + } + } + } + + #[test] + fn test_datetime_date_add() { + let f = DateAddFunction; + + let dates = vec![Some(123), None, Some(42), None]; + // Intervals in months + let intervals = vec![1, 2, 3, 1]; + let results = [Some(2678400123), None, Some(7776000042), None]; + + let date_vector = DateTimeVector::from(dates.clone()); + let interval_vector = IntervalYearMonthVector::from_vec(intervals); + let args: Vec = vec![Arc::new(date_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in dates.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::DateTime(date) => { + assert_eq!(date.val(), result.unwrap()); + } + _ => unreachable!(), + } + } + } +} diff --git a/src/common/function/src/scalars/date/date_sub.rs b/src/common/function/src/scalars/date/date_sub.rs new file mode 100644 index 000000000000..15660850f558 --- /dev/null +++ b/src/common/function/src/scalars/date/date_sub.rs @@ -0,0 +1,292 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt; + +use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu}; +use common_query::prelude::Signature; +use datatypes::data_type::DataType; +use datatypes::prelude::ConcreteDataType; +use datatypes::value::ValueRef; +use datatypes::vectors::VectorRef; +use snafu::ensure; + +use crate::helper; +use crate::scalars::function::{Function, FunctionContext}; + +/// A function subtracts an interval value to Timestamp, Date or DateTime, and return the result. +#[derive(Clone, Debug, Default)] +pub struct DateSubFunction; + +const NAME: &str = "date_sub"; + +impl Function for DateSubFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + Ok(input_types[0].clone()) + } + + fn signature(&self) -> Signature { + helper::one_of_sigs2( + vec![ + ConcreteDataType::date_datatype(), + ConcreteDataType::datetime_datatype(), + ConcreteDataType::timestamp_second_datatype(), + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_microsecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ], + vec![ + ConcreteDataType::interval_month_day_nano_datatype(), + ConcreteDataType::interval_year_month_datatype(), + ConcreteDataType::interval_day_time_datatype(), + ], + ) + } + + fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result { + ensure!( + columns.len() == 2, + InvalidFuncArgsSnafu { + err_msg: format!( + "The length of the args is not correct, expect 2, have: {}", + columns.len() + ), + } + ); + + let left = &columns[0]; + let right = &columns[1]; + + let size = left.len(); + let left_datatype = columns[0].data_type(); + + match left_datatype { + ConcreteDataType::Timestamp(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let ts = left.get(i).as_timestamp(); + let interval = right.get(i).as_interval(); + + let new_ts = match (ts, interval) { + (Some(ts), Some(interval)) => ts.sub_interval(interval), + _ => ts, + }; + + result.push_value_ref(ValueRef::from(new_ts)); + } + + Ok(result.to_vector()) + } + ConcreteDataType::Date(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let date = left.get(i).as_date(); + let interval = right.get(i).as_interval(); + let new_date = match (date, interval) { + (Some(date), Some(interval)) => date.sub_interval(interval), + _ => date, + }; + + result.push_value_ref(ValueRef::from(new_date)); + } + + Ok(result.to_vector()) + } + ConcreteDataType::DateTime(_) => { + let mut result = left_datatype.create_mutable_vector(size); + for i in 0..size { + let datetime = left.get(i).as_datetime(); + let interval = right.get(i).as_interval(); + let new_datetime = match (datetime, interval) { + (Some(datetime), Some(interval)) => datetime.sub_interval(interval), + _ => datetime, + }; + + result.push_value_ref(ValueRef::from(new_datetime)); + } + + Ok(result.to_vector()) + } + _ => UnsupportedInputDataTypeSnafu { + function: NAME, + datatypes: columns.iter().map(|c| c.data_type()).collect::>(), + } + .fail(), + } + } +} + +impl fmt::Display for DateSubFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "DATE_SUB") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_query::prelude::{TypeSignature, Volatility}; + use datatypes::prelude::ConcreteDataType; + use datatypes::value::Value; + use datatypes::vectors::{ + DateTimeVector, DateVector, IntervalDayTimeVector, IntervalYearMonthVector, + TimestampSecondVector, + }; + + use super::{DateSubFunction, *}; + use crate::scalars::Function; + + #[test] + fn test_date_sub_misc() { + let f = DateSubFunction; + assert_eq!("date_sub", f.name()); + assert_eq!( + ConcreteDataType::timestamp_microsecond_datatype(), + f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()]) + .unwrap() + ); + assert_eq!( + ConcreteDataType::timestamp_second_datatype(), + f.return_type(&[ConcreteDataType::timestamp_second_datatype()]) + .unwrap() + ); + assert_eq!( + ConcreteDataType::date_datatype(), + f.return_type(&[ConcreteDataType::date_datatype()]).unwrap() + ); + assert_eq!( + ConcreteDataType::datetime_datatype(), + f.return_type(&[ConcreteDataType::datetime_datatype()]) + .unwrap() + ); + assert!(matches!(f.signature(), + Signature { + type_signature: TypeSignature::OneOf(sigs), + volatility: Volatility::Immutable + } if sigs.len() == 18)); + } + + #[test] + fn test_timestamp_date_sub() { + let f = DateSubFunction; + + let times = vec![Some(123), None, Some(42), None]; + // Intervals in milliseconds + let intervals = vec![1000, 2000, 3000, 1000]; + let results = [Some(122), None, Some(39), None]; + + let time_vector = TimestampSecondVector::from(times.clone()); + let interval_vector = IntervalDayTimeVector::from_vec(intervals); + let args: Vec = vec![Arc::new(time_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in times.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::Timestamp(ts) => { + assert_eq!(ts.value(), result.unwrap()); + } + _ => unreachable!(), + } + } + } + + #[test] + fn test_date_date_sub() { + let f = DateSubFunction; + let days_per_month = 30; + + let dates = vec![ + Some(123 * days_per_month), + None, + Some(42 * days_per_month), + None, + ]; + // Intervals in months + let intervals = vec![1, 2, 3, 1]; + let results = [Some(3659), None, Some(1168), None]; + + let date_vector = DateVector::from(dates.clone()); + let interval_vector = IntervalYearMonthVector::from_vec(intervals); + let args: Vec = vec![Arc::new(date_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in dates.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::Date(date) => { + assert_eq!(date.val(), result.unwrap()); + } + _ => unreachable!(), + } + } + } + + #[test] + fn test_datetime_date_sub() { + let f = DateSubFunction; + let millis_per_month = 3600 * 24 * 30 * 1000; + + let dates = vec![ + Some(123 * millis_per_month), + None, + Some(42 * millis_per_month), + None, + ]; + // Intervals in months + let intervals = vec![1, 2, 3, 1]; + let results = [Some(316137600000), None, Some(100915200000), None]; + + let date_vector = DateTimeVector::from(dates.clone()); + let interval_vector = IntervalYearMonthVector::from_vec(intervals); + let args: Vec = vec![Arc::new(date_vector), Arc::new(interval_vector)]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + + assert_eq!(4, vector.len()); + for (i, _t) in dates.iter().enumerate() { + let v = vector.get(i); + let result = results.get(i).unwrap(); + + if result.is_none() { + assert_eq!(Value::Null, v); + continue; + } + match v { + Value::DateTime(date) => { + assert_eq!(date.val(), result.unwrap()); + } + _ => unreachable!(), + } + } + } +} diff --git a/src/common/function/src/scalars/function_registry.rs b/src/common/function/src/scalars/function_registry.rs index 1fd536982e35..5d6751df7503 100644 --- a/src/common/function/src/scalars/function_registry.rs +++ b/src/common/function/src/scalars/function_registry.rs @@ -19,6 +19,7 @@ use std::sync::{Arc, RwLock}; use once_cell::sync::Lazy; use crate::scalars::aggregate::{AggregateFunctionMetaRef, AggregateFunctions}; +use crate::scalars::date::DateFunction; use crate::scalars::function::FunctionRef; use crate::scalars::math::MathFunction; use crate::scalars::numpy::NumpyFunction; @@ -75,6 +76,7 @@ pub static FUNCTION_REGISTRY: Lazy> = Lazy::new(|| { MathFunction::register(&function_registry); NumpyFunction::register(&function_registry); TimestampFunction::register(&function_registry); + DateFunction::register(&function_registry); AggregateFunctions::register(&function_registry); diff --git a/src/common/time/src/date.rs b/src/common/time/src/date.rs index b94c483e3081..7ff19145430a 100644 --- a/src/common/time/src/date.rs +++ b/src/common/time/src/date.rs @@ -15,12 +15,13 @@ use std::fmt::{Display, Formatter}; use std::str::FromStr; -use chrono::{Datelike, NaiveDate}; +use chrono::{Datelike, Days, Months, NaiveDate}; use serde::{Deserialize, Serialize}; use serde_json::Value; use snafu::ResultExt; use crate::error::{Error, ParseDateStrSnafu, Result}; +use crate::interval::Interval; const UNIX_EPOCH_FROM_CE: i32 = 719_163; @@ -86,6 +87,32 @@ impl Date { pub fn to_secs(&self) -> i64 { (self.0 as i64) * 24 * 3600 } + + /// Adds given Interval to the current date. + /// Returns None if the resulting date would be out of range. + pub fn add_interval(&self, interval: Interval) -> Option { + let naive_date = self.to_chrono_date()?; + + let (months, days, _) = interval.to_month_day_nano(); + + naive_date + .checked_add_months(Months::new(months as u32))? + .checked_add_days(Days::new(days as u64)) + .map(Into::into) + } + + /// Subtracts given Interval to the current date. + /// Returns None if the resulting date would be out of range. + pub fn sub_interval(&self, interval: Interval) -> Option { + let naive_date = self.to_chrono_date()?; + + let (months, days, _) = interval.to_month_day_nano(); + + naive_date + .checked_sub_months(Months::new(months as u32))? + .checked_sub_days(Days::new(days as u64)) + .map(Into::into) + } } #[cfg(test)] @@ -124,6 +151,18 @@ mod tests { assert_eq!(now, Date::from_str(&now).unwrap().to_string()); } + #[test] + fn test_add_sub_interval() { + let date = Date::new(1000); + + let interval = Interval::from_year_month(3); + + let new_date = date.add_interval(interval).unwrap(); + assert_eq!(new_date.val(), 1091); + + assert_eq!(date, new_date.sub_interval(interval).unwrap()); + } + #[test] pub fn test_min_max() { let mut date = Date::from_str("9999-12-31").unwrap(); diff --git a/src/common/time/src/datetime.rs b/src/common/time/src/datetime.rs index 7f1517523072..3a2274cae7fb 100644 --- a/src/common/time/src/datetime.rs +++ b/src/common/time/src/datetime.rs @@ -14,14 +14,15 @@ use std::fmt::{Display, Formatter}; use std::str::FromStr; +use std::time::Duration; -use chrono::{LocalResult, NaiveDateTime, TimeZone as ChronoTimeZone, Utc}; +use chrono::{Days, LocalResult, Months, NaiveDateTime, TimeZone as ChronoTimeZone, Utc}; use serde::{Deserialize, Serialize}; use crate::error::{Error, InvalidDateStrSnafu, Result}; use crate::timezone::TimeZone; use crate::util::{format_utc_datetime, local_datetime_to_utc}; -use crate::Date; +use crate::{Date, Interval}; const DATETIME_FORMAT: &str = "%F %T"; const DATETIME_FORMAT_WITH_TZ: &str = "%F %T%z"; @@ -117,6 +118,33 @@ impl DateTime { None => Utc.from_utc_datetime(&v).naive_local(), }) } + /// Adds given Interval to the current datetime. + /// Returns None if the resulting datetime would be out of range. + pub fn add_interval(&self, interval: Interval) -> Option { + let naive_datetime = self.to_chrono_datetime()?; + let (months, days, nsecs) = interval.to_month_day_nano(); + + let naive_datetime = naive_datetime + .checked_add_months(Months::new(months as u32))? + .checked_add_days(Days::new(days as u64))? + + Duration::from_nanos(nsecs as u64); + + Some(naive_datetime.into()) + } + + /// Subtracts given Interval to the current datetime. + /// Returns None if the resulting datetime would be out of range. + pub fn sub_interval(&self, interval: Interval) -> Option { + let naive_datetime = self.to_chrono_datetime()?; + let (months, days, nsecs) = interval.to_month_day_nano(); + + let naive_datetime = naive_datetime + .checked_sub_months(Months::new(months as u32))? + .checked_sub_days(Days::new(days as u64))? + - Duration::from_nanos(nsecs as u64); + + Some(naive_datetime.into()) + } /// Convert to [common_time::date]. pub fn to_date(&self) -> Option { @@ -152,6 +180,18 @@ mod tests { assert_eq!(42, d.val()); } + #[test] + fn test_add_sub_interval() { + let datetime = DateTime::new(1000); + + let interval = Interval::from_day_time(1, 200); + + let new_datetime = datetime.add_interval(interval).unwrap(); + assert_eq!(new_datetime.val(), 1000 + 3600 * 24 * 1000 + 200); + + assert_eq!(datetime, new_datetime.sub_interval(interval).unwrap()); + } + #[test] fn test_parse_local_date_time() { std::env::set_var("TZ", "Asia/Shanghai"); diff --git a/src/common/time/src/interval.rs b/src/common/time/src/interval.rs index a7c184c0ec47..716602b8e2a8 100644 --- a/src/common/time/src/interval.rs +++ b/src/common/time/src/interval.rs @@ -20,6 +20,10 @@ use std::hash::{Hash, Hasher}; use arrow::datatypes::IntervalUnit as ArrowIntervalUnit; use serde::{Deserialize, Serialize}; use serde_json::Value; +use snafu::ResultExt; + +use crate::duration::Duration; +use crate::error::{Result, TimestampOverflowSnafu}; #[derive( Debug, Default, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, @@ -63,7 +67,7 @@ impl From for IntervalUnit { /// month-day-nano, which will be stored in the following format. /// Interval data format: /// | months | days | nsecs | -/// | 4bytes | 4bytes | 8bytes | +/// | 4bytes | 4bytes | 8bytes | #[derive(Debug, Clone, Default, Copy, Serialize, Deserialize)] pub struct Interval { months: i32, @@ -114,6 +118,14 @@ impl Interval { } } + pub fn to_duration(&self) -> Result { + Ok(Duration::new_nanosecond( + self.to_nanosecond() + .try_into() + .context(TimestampOverflowSnafu)?, + )) + } + /// Return a tuple(months, days, nanoseconds) from the interval. pub fn to_month_day_nano(&self) -> (i32, i32, i64) { (self.months, self.days, self.nsecs) @@ -558,6 +570,7 @@ mod tests { use std::collections::HashMap; use super::*; + use crate::timestamp::TimeUnit; #[test] fn test_from_year_month() { @@ -572,6 +585,21 @@ mod tests { assert_eq!(interval.nsecs, 2_000_000); } + #[test] + fn test_to_duration() { + let interval = Interval::from_day_time(1, 2); + + let duration = interval.to_duration().unwrap(); + assert_eq!(86400002000000, duration.value()); + assert_eq!(TimeUnit::Nanosecond, duration.unit()); + + let interval = Interval::from_year_month(12); + + let duration = interval.to_duration().unwrap(); + assert_eq!(31104000000000000, duration.value()); + assert_eq!(TimeUnit::Nanosecond, duration.unit()); + } + #[test] fn test_interval_is_positive() { let interval = Interval::from_year_month(1); diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 898e1c7b39c4..e8cf29a68737 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -21,15 +21,16 @@ use std::time::Duration; use arrow::datatypes::TimeUnit as ArrowTimeUnit; use chrono::{ - DateTime, LocalResult, NaiveDate, NaiveDateTime, NaiveTime, TimeZone as ChronoTimeZone, Utc, + DateTime, Days, LocalResult, Months, NaiveDate, NaiveDateTime, NaiveTime, + TimeZone as ChronoTimeZone, Utc, }; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; -use crate::error; use crate::error::{ArithmeticOverflowSnafu, Error, ParseTimestampSnafu, TimestampOverflowSnafu}; use crate::timezone::TimeZone; use crate::util::{div_ceil, format_utc_datetime, local_datetime_to_utc}; +use crate::{error, Interval}; /// Timestamp represents the value of units(seconds/milliseconds/microseconds/nanoseconds) elapsed /// since UNIX epoch. The valid value range of [Timestamp] depends on it's unit (all in UTC time zone): @@ -104,6 +105,78 @@ impl Timestamp { }) } + /// Adds a duration to timestamp. + /// # Note + /// The result time unit remains unchanged even if `duration` has a different unit with `self`. + /// For example, a timestamp with value 1 and time unit second, subtracted by 1 millisecond + /// and the result is still 1 second. + pub fn add_duration(&self, duration: Duration) -> error::Result { + let duration: i64 = match self.unit { + TimeUnit::Second => { + i64::try_from(duration.as_secs()).context(TimestampOverflowSnafu)? + } + TimeUnit::Millisecond => { + i64::try_from(duration.as_millis()).context(TimestampOverflowSnafu)? + } + TimeUnit::Microsecond => { + i64::try_from(duration.as_micros()).context(TimestampOverflowSnafu)? + } + TimeUnit::Nanosecond => { + i64::try_from(duration.as_nanos()).context(TimestampOverflowSnafu)? + } + }; + + let value = self + .value + .checked_add(duration) + .with_context(|| ArithmeticOverflowSnafu { + msg: format!( + "Try to add timestamp: {:?} with duration: {:?}", + self, duration + ), + })?; + Ok(Timestamp { + value, + unit: self.unit, + }) + } + + /// Adds given Interval to the current timestamp. + /// Returns None if the resulting timestamp would be out of range. + pub fn add_interval(&self, interval: Interval) -> Option { + let naive_datetime = self.to_chrono_datetime()?; + let (months, days, nsecs) = interval.to_month_day_nano(); + + let naive_datetime = naive_datetime + .checked_add_months(Months::new(months as u32))? + .checked_add_days(Days::new(days as u64))? + + Duration::from_nanos(nsecs as u64); + + match Timestamp::from_chrono_datetime(naive_datetime) { + // Have to convert the new timestamp by the current unit. + Some(ts) => ts.convert_to(self.unit), + None => None, + } + } + + /// Subtracts given Interval to the current timestamp. + /// Returns None if the resulting timestamp would be out of range. + pub fn sub_interval(&self, interval: Interval) -> Option { + let naive_datetime = self.to_chrono_datetime()?; + let (months, days, nsecs) = interval.to_month_day_nano(); + + let naive_datetime = naive_datetime + .checked_sub_months(Months::new(months as u32))? + .checked_sub_days(Days::new(days as u64))? + - Duration::from_nanos(nsecs as u64); + + match Timestamp::from_chrono_datetime(naive_datetime) { + // Have to convert the new timestamp by the current unit. + Some(ts) => ts.convert_to(self.unit), + None => None, + } + } + /// Subtracts current timestamp with another timestamp, yielding a duration. pub fn sub(&self, rhs: &Self) -> Option { let lhs = self.to_chrono_datetime()?; @@ -543,6 +616,19 @@ mod tests { Timestamp::new(value, unit) } + #[test] + fn test_add_sub_interval() { + let ts = Timestamp::new(1000, TimeUnit::Millisecond); + + let interval = Interval::from_day_time(1, 200); + + let new_ts = ts.add_interval(interval).unwrap(); + assert_eq!(new_ts.unit(), TimeUnit::Millisecond); + assert_eq!(new_ts.value(), 1000 + 3600 * 24 * 1000 + 200); + + assert_eq!(ts, new_ts.sub_interval(interval).unwrap()); + } + #[test] fn test_timestamp_reflexivity() { for _ in 0..1000 { @@ -1006,6 +1092,33 @@ mod tests { assert_eq!(TimeUnit::Second, res.unit); } + #[test] + fn test_timestamp_add() { + let res = Timestamp::new(1, TimeUnit::Second) + .add_duration(Duration::from_secs(1)) + .unwrap(); + assert_eq!(2, res.value); + assert_eq!(TimeUnit::Second, res.unit); + + let res = Timestamp::new(0, TimeUnit::Second) + .add_duration(Duration::from_secs(1)) + .unwrap(); + assert_eq!(1, res.value); + assert_eq!(TimeUnit::Second, res.unit); + + let res = Timestamp::new(1, TimeUnit::Second) + .add_duration(Duration::from_millis(1)) + .unwrap(); + assert_eq!(1, res.value); + assert_eq!(TimeUnit::Second, res.unit); + + let res = Timestamp::new(100, TimeUnit::Second) + .add_duration(Duration::from_millis(1000)) + .unwrap(); + assert_eq!(101, res.value); + assert_eq!(TimeUnit::Second, res.unit); + } + #[test] fn test_parse_in_time_zone() { std::env::set_var("TZ", "Asia/Shanghai"); diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 23135a6e44a8..c198fde9a7b1 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -210,6 +210,14 @@ impl Value { } } + /// Cast Value to Interval. Return None if value is not a valid interval data type. + pub fn as_interval(&self) -> Option { + match self { + Value::Interval(i) => Some(*i), + _ => None, + } + } + /// Cast Value to Date. Return None if value is not a valid date data type. pub fn as_date(&self) -> Option { match self { diff --git a/tests/cases/standalone/common/function/date.result b/tests/cases/standalone/common/function/date.result new file mode 100644 index 000000000000..122866efac4d --- /dev/null +++ b/tests/cases/standalone/common/function/date.result @@ -0,0 +1,164 @@ +SELECT date_add('2023-12-06 07:39:46.222'::TIMESTAMP_MS, INTERVAL '5 day'); + ++----------------------------------------------------------------------------------------+ +| date_add(Utf8("2023-12-06 07:39:46.222"),IntervalMonthDayNano("92233720368547758080")) | ++----------------------------------------------------------------------------------------+ +| 2023-12-11T07:39:46.222 | ++----------------------------------------------------------------------------------------+ + +SELECT date_add('2023-12-06 07:39:46.222'::TIMESTAMP_MS, '5 day'); + ++---------------------------------------------------------+ +| date_add(Utf8("2023-12-06 07:39:46.222"),Utf8("5 day")) | ++---------------------------------------------------------+ +| 2023-12-11T07:39:46.222 | ++---------------------------------------------------------+ + +SELECT date_add('2023-12-06'::DATE, INTERVAL '3 month 5 day'); + ++-------------------------------------------------------------------------------------+ +| date_add(Utf8("2023-12-06"),IntervalMonthDayNano("237684487635026733149179609088")) | ++-------------------------------------------------------------------------------------+ +| 2024-03-11 | ++-------------------------------------------------------------------------------------+ + +SELECT date_add('2023-12-06'::DATE, '3 month 5 day'); + ++----------------------------------------------------+ +| date_add(Utf8("2023-12-06"),Utf8("3 month 5 day")) | ++----------------------------------------------------+ +| 2024-03-11 | ++----------------------------------------------------+ + +SELECT date_sub('2023-12-06 07:39:46.222'::TIMESTAMP_MS, INTERVAL '5 day'); + ++----------------------------------------------------------------------------------------+ +| date_sub(Utf8("2023-12-06 07:39:46.222"),IntervalMonthDayNano("92233720368547758080")) | ++----------------------------------------------------------------------------------------+ +| 2023-12-01T07:39:46.222 | ++----------------------------------------------------------------------------------------+ + +SELECT date_sub('2023-12-06 07:39:46.222'::TIMESTAMP_MS, '5 day'); + ++---------------------------------------------------------+ +| date_sub(Utf8("2023-12-06 07:39:46.222"),Utf8("5 day")) | ++---------------------------------------------------------+ +| 2023-12-01T07:39:46.222 | ++---------------------------------------------------------+ + +SELECT date_sub('2023-12-06'::DATE, INTERVAL '3 month 5 day'); + ++-------------------------------------------------------------------------------------+ +| date_sub(Utf8("2023-12-06"),IntervalMonthDayNano("237684487635026733149179609088")) | ++-------------------------------------------------------------------------------------+ +| 2023-09-01 | ++-------------------------------------------------------------------------------------+ + +SELECT date_sub('2023-12-06'::DATE, '3 month 5 day'); + ++----------------------------------------------------+ +| date_sub(Utf8("2023-12-06"),Utf8("3 month 5 day")) | ++----------------------------------------------------+ +| 2023-09-01 | ++----------------------------------------------------+ + +CREATE TABLE dates(d DATE, ts timestamp time index); + +Affected Rows: 0 + +INSERT INTO dates VALUES ('1992-01-01'::DATE, 1); + +Affected Rows: 1 + +INSERT INTO dates VALUES ('1993-12-30'::DATE, 2); + +Affected Rows: 1 + +INSERT INTO dates VALUES ('2023-12-06'::DATE, 3); + +Affected Rows: 1 + +SELECT date_add(d, INTERVAL '1 year 2 month 3 day') from dates; + ++---------------------------------------------------------------------------+ +| date_add(dates.d,IntervalMonthDayNano("1109194275255040958530743959552")) | ++---------------------------------------------------------------------------+ +| 1993-03-04 | +| 1995-03-03 | +| 2025-02-09 | ++---------------------------------------------------------------------------+ + +SELECT date_add(d, '1 year 2 month 3 day') from dates; + ++------------------------------------------------+ +| date_add(dates.d,Utf8("1 year 2 month 3 day")) | ++------------------------------------------------+ +| 1993-03-04 | +| 1995-03-03 | +| 2025-02-09 | ++------------------------------------------------+ + +SELECT date_add(ts, INTERVAL '1 year 2 month 3 day') from dates; + ++----------------------------------------------------------------------------+ +| date_add(dates.ts,IntervalMonthDayNano("1109194275255040958530743959552")) | ++----------------------------------------------------------------------------+ +| 1971-03-04T00:00:00.001 | +| 1971-03-04T00:00:00.002 | +| 1971-03-04T00:00:00.003 | ++----------------------------------------------------------------------------+ + +SELECT date_add(ts, '1 year 2 month 3 day') from dates; + ++-------------------------------------------------+ +| date_add(dates.ts,Utf8("1 year 2 month 3 day")) | ++-------------------------------------------------+ +| 1971-03-04T00:00:00.001 | +| 1971-03-04T00:00:00.002 | +| 1971-03-04T00:00:00.003 | ++-------------------------------------------------+ + +SELECT date_sub(d, INTERVAL '1 year 2 month 3 day') from dates; + ++---------------------------------------------------------------------------+ +| date_sub(dates.d,IntervalMonthDayNano("1109194275255040958530743959552")) | ++---------------------------------------------------------------------------+ +| 1990-10-29 | +| 1992-10-27 | +| 2022-10-03 | ++---------------------------------------------------------------------------+ + +SELECT date_sub(d, '1 year 2 month 3 day') from dates; + ++------------------------------------------------+ +| date_sub(dates.d,Utf8("1 year 2 month 3 day")) | ++------------------------------------------------+ +| 1990-10-29 | +| 1992-10-27 | +| 2022-10-03 | ++------------------------------------------------+ + +SELECT date_sub(ts, INTERVAL '1 year 2 month 3 day') from dates; + ++----------------------------------------------------------------------------+ +| date_sub(dates.ts,IntervalMonthDayNano("1109194275255040958530743959552")) | ++----------------------------------------------------------------------------+ +| 1968-10-29T00:00:00.001 | +| 1968-10-29T00:00:00.002 | +| 1968-10-29T00:00:00.003 | ++----------------------------------------------------------------------------+ + +SELECT date_sub(ts, '1 year 2 month 3 day') from dates; + ++-------------------------------------------------+ +| date_sub(dates.ts,Utf8("1 year 2 month 3 day")) | ++-------------------------------------------------+ +| 1968-10-29T00:00:00.001 | +| 1968-10-29T00:00:00.002 | +| 1968-10-29T00:00:00.003 | ++-------------------------------------------------+ + +DROP TABLE dates; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/date.sql b/tests/cases/standalone/common/function/date.sql new file mode 100644 index 000000000000..a5f962749975 --- /dev/null +++ b/tests/cases/standalone/common/function/date.sql @@ -0,0 +1,42 @@ +SELECT date_add('2023-12-06 07:39:46.222'::TIMESTAMP_MS, INTERVAL '5 day'); + +SELECT date_add('2023-12-06 07:39:46.222'::TIMESTAMP_MS, '5 day'); + +SELECT date_add('2023-12-06'::DATE, INTERVAL '3 month 5 day'); + +SELECT date_add('2023-12-06'::DATE, '3 month 5 day'); + +SELECT date_sub('2023-12-06 07:39:46.222'::TIMESTAMP_MS, INTERVAL '5 day'); + +SELECT date_sub('2023-12-06 07:39:46.222'::TIMESTAMP_MS, '5 day'); + +SELECT date_sub('2023-12-06'::DATE, INTERVAL '3 month 5 day'); + +SELECT date_sub('2023-12-06'::DATE, '3 month 5 day'); + + +CREATE TABLE dates(d DATE, ts timestamp time index); + +INSERT INTO dates VALUES ('1992-01-01'::DATE, 1); + +INSERT INTO dates VALUES ('1993-12-30'::DATE, 2); + +INSERT INTO dates VALUES ('2023-12-06'::DATE, 3); + +SELECT date_add(d, INTERVAL '1 year 2 month 3 day') from dates; + +SELECT date_add(d, '1 year 2 month 3 day') from dates; + +SELECT date_add(ts, INTERVAL '1 year 2 month 3 day') from dates; + +SELECT date_add(ts, '1 year 2 month 3 day') from dates; + +SELECT date_sub(d, INTERVAL '1 year 2 month 3 day') from dates; + +SELECT date_sub(d, '1 year 2 month 3 day') from dates; + +SELECT date_sub(ts, INTERVAL '1 year 2 month 3 day') from dates; + +SELECT date_sub(ts, '1 year 2 month 3 day') from dates; + +DROP TABLE dates;