rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-01-16 02:02:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: impl date_add/date_sub functions (#2881)

Browse Source 
* feat: adds date_add and date_sub function

* test: add date function

* fix: adds interval to date returns wrong result

* fix: header

* fix: typo

* fix: timestamp resolution

* fix: capacity

* chore: apply suggestion

* fix: wrong behavior when adding intervals to timestamp, date and datetime

* chore: remove unused error

* test: refactor and add some tests
This commit is contained in:
dennis zhuang
2023-12-07 11:02:15 +08:00
committed by GitHub
parent ea80570cb1
commit 43f32f4499
14 changed files with 1076 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/common/function/src/helper.rs Normal file
View File

@@ -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<ConcreteDataType>, args2: Vec<ConcreteDataType>) -> 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)
}

2
src/common/function/src/lib.rs
View File

@@ -13,3 +13,5 @@
// limitations under the License.
pub mod scalars;
pub mod helper;

1
src/common/function/src/scalars.rs
View File

@@ -13,6 +13,7 @@
// limitations under the License.
pub mod aggregate;
mod date;
pub mod expression;
pub mod function;
pub mod function_registry;

31
src/common/function/src/scalars/date.rs Normal file
View File

@@ -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));
}
}

279
src/common/function/src/scalars/date/date_add.rs Normal file
View File

@@ -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<ConcreteDataType> {
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<VectorRef> {
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::<Vec<_>>(),
}
.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<VectorRef> = 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<VectorRef> = 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<VectorRef> = 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!(),
}
}
}
}

292
src/common/function/src/scalars/date/date_sub.rs Normal file
View File

@@ -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<ConcreteDataType> {
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<VectorRef> {
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::<Vec<_>>(),
}
.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<VectorRef> = 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<VectorRef> = 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<VectorRef> = 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!(),
}
}
}
}

2
src/common/function/src/scalars/function_registry.rs
View File

@@ -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<Arc<FunctionRegistry>> = Lazy::new(|| {
MathFunction::register(&function_registry);
NumpyFunction::register(&function_registry);
TimestampFunction::register(&function_registry);
DateFunction::register(&function_registry);
AggregateFunctions::register(&function_registry);

41
src/common/time/src/date.rs
View File

@@ -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<Date> {
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<Date> {
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();

44
src/common/time/src/datetime.rs
View File

@@ -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<Self> {
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<Self> {
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<Date> {
@@ -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");

30
src/common/time/src/interval.rs
View File

@@ -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<ArrowIntervalUnit> 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<Duration> {
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);

117
src/common/time/src/timestamp.rs
View File

@@ -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<Self> {
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<Timestamp> {
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<Timestamp> {
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<chrono::Duration> {
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");

8
src/datatypes/src/value.rs
View File

@@ -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<Interval> {
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<Date> {
match self {

164
tests/cases/standalone/common/function/date.result Normal file
View File

@@ -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

42
tests/cases/standalone/common/function/date.sql Normal file
View File

@@ -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;