Skip to content

Commit

Permalink
support more units for extract
Browse files Browse the repository at this point in the history 
Signed-off-by: Runji Wang <[email protected]>
  • Loading branch information
@wangrunji0408
wangrunji0408 committed Mar 28, 2023
1 parent 8f28db3 commit 795a19f
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 79 deletions.
90 changes: 62 additions & 28 deletions src/expr/src/vector_op/extract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,35 +12,45 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use chrono::{Datelike, Timelike};
use chrono::{Datelike, NaiveDate, NaiveTime, Timelike};
use risingwave_common::types::{Date, Decimal, Time, Timestamp};
use risingwave_expr_macro::function;

use crate::{ExprError, Result};

fn extract_time<T>(time: T, unit: &str) -> Option<Decimal>
where
T: Timelike,
{
fn extract_time(time: NaiveTime, unit: &str) -> Option<Decimal> {
Some(match unit {
"HOUR" => time.hour().into(),
"MINUTE" => time.minute().into(),
"SECOND" => time.second().into(),
"HOUR" | "hour" => time.hour().into(),
"MINUTE" | "minute" => time.minute().into(),
"SECOND" | "second" => time.second().into(),
"MILLISECOND" | "millisecond" => {
(time.second() * 1_000 + time.nanosecond() / 1_000_000).into()
}
"MICROSECOND" | "microsecond" => {
(time.second() * 1_000_000 + time.nanosecond() / 1_000).into()
}
"EPOCH" | "epoch" => (time.num_seconds_from_midnight() as f64
+ time.nanosecond() as f64 / 1_000_000_000.0)
.into(),
_ => return None,
})
}

fn extract_date<T>(date: T, unit: &str) -> Option<Decimal>
where
T: Datelike,
{
fn extract_date(date: NaiveDate, unit: &str) -> Option<Decimal> {
Some(match unit {
"DAY" => date.day().into(),
"MONTH" => date.month().into(),
"YEAR" => date.year().into(),
// Sun = 0 and Sat = 6
"DOW" => date.weekday().num_days_from_sunday().into(),
"DOY" => date.ordinal().into(),
"MILLENNIUM" | "millennium" => ((date.year() - 1) / 1000 + 1).into(),
"CENTURY" | "century" => ((date.year() - 1) / 100 + 1).into(),
"DECADE" | "decade" => (date.year() / 10).into(),
"YEAR" | "year" => date.year().into(),
"ISOYEAR" | "isoyear" => date.iso_week().year().into(),
"QUARTER" | "quarter" => ((date.month() - 1) / 3 + 1).into(),
"MONTH" | "month" => date.month().into(),
"WEEK" | "week" => date.iso_week().week().into(),
"DAY" | "day" => date.day().into(),
"DOY" | "doy" => date.ordinal().into(),
"DOW" | "dow" => date.weekday().num_days_from_sunday().into(),
"ISODOW" | "isodow" => date.weekday().number_from_monday().into(),
"EPOCH" | "epoch" => date.and_time(NaiveTime::default()).timestamp().into(),
_ => return None,
})
}
Expand All @@ -59,10 +69,13 @@ pub fn extract_from_date(unit: &str, date: Date) -> Result<Decimal> {

#[function("extract(varchar, timestamp) -> decimal")]
pub fn extract_from_timestamp(unit: &str, timestamp: Timestamp) -> Result<Decimal> {
let time = timestamp.0;

extract_date(time, unit)
.or_else(|| extract_time(time, unit))
if matches!(unit, "EPOCH" | "epoch") {
let epoch =
timestamp.0.timestamp() as f64 + timestamp.0.nanosecond() as f64 / 1_000_000_000.0;
return Ok(epoch.into());
}
extract_date(timestamp.0.date(), unit)
.or_else(|| extract_time(timestamp.0.time(), unit))
.ok_or_else(|| invalid_unit("timestamp unit", unit))
}

Expand Down Expand Up @@ -94,15 +107,36 @@ mod tests {
assert_eq!(extract_from_date("YEAR", date).unwrap(), 2021.into());
assert_eq!(extract_from_date("DOW", date).unwrap(), 1.into());
assert_eq!(extract_from_date("DOY", date).unwrap(), 326.into());
assert_eq!(extract_from_date("MILLENNIUM", date).unwrap(), 3.into());
assert_eq!(extract_from_date("CENTURY", date).unwrap(), 21.into());
assert_eq!(extract_from_date("DECADE", date).unwrap(), 202.into());
assert_eq!(extract_from_date("ISOYEAR", date).unwrap(), 2021.into());
assert_eq!(extract_from_date("QUARTER", date).unwrap(), 4.into());
assert_eq!(extract_from_date("WEEK", date).unwrap(), 47.into());
assert_eq!(extract_from_date("ISODOW", date).unwrap(), 1.into());
assert_eq!(extract_from_date("EPOCH", date).unwrap(), 1637539200.into());
}

#[test]
fn test_time() {
let time = Timestamp::new(
NaiveDateTime::parse_from_str("2021-11-22 12:4:2", "%Y-%m-%d %H:%M:%S").unwrap(),
fn test_timestamp() {
let ts = Timestamp::new(
NaiveDateTime::parse_from_str("2021-11-22 12:4:2.575401000", "%Y-%m-%d %H:%M:%S%.f")
.unwrap(),
);
assert_eq!(extract_from_timestamp("HOUR", ts).unwrap(), 12.into());
assert_eq!(extract_from_timestamp("MINUTE", ts).unwrap(), 4.into());
assert_eq!(extract_from_timestamp("SECOND", ts).unwrap(), 2.into());
assert_eq!(
extract_from_timestamp("MILLISECOND", ts).unwrap(),
2575.into()
);
assert_eq!(
extract_from_timestamp("MICROSECOND", ts).unwrap(),
2575401.into()
);
assert_eq!(
extract_from_timestamp("EPOCH", ts).unwrap(),
1637582642.575401.into()
);
assert_eq!(extract_from_timestamp("HOUR", time).unwrap(), 12.into());
assert_eq!(extract_from_timestamp("MINUTE", time).unwrap(), 4.into());
assert_eq!(extract_from_timestamp("SECOND", time).unwrap(), 2.into());
}
}
86 changes: 43 additions & 43 deletions src/tests/regress/data/sql/date.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,22 +223,22 @@ SELECT f1 - date '2000-01-01' AS "Days From 2K" FROM DATE_TBL;
--
-- test extract!
--
--@ SELECT f1 as "date",
--@ date_part('year', f1) AS year,
--@ date_part('month', f1) AS month,
--@ date_part('day', f1) AS day,
--@ date_part('quarter', f1) AS quarter,
--@ date_part('decade', f1) AS decade,
--@ date_part('century', f1) AS century,
--@ date_part('millennium', f1) AS millennium,
--@ date_part('isoyear', f1) AS isoyear,
--@ date_part('week', f1) AS week,
--@ date_part('dow', f1) AS dow,
--@ date_part('isodow', f1) AS isodow,
--@ date_part('doy', f1) AS doy,
--@ date_part('julian', f1) AS julian,
--@ date_part('epoch', f1) AS epoch
--@ FROM date_tbl;
SELECT f1 as "date",
date_part('year', f1) AS year,
date_part('month', f1) AS month,
date_part('day', f1) AS day,
date_part('quarter', f1) AS quarter,
date_part('decade', f1) AS decade,
date_part('century', f1) AS century,
date_part('millennium', f1) AS millennium,
date_part('isoyear', f1) AS isoyear,
date_part('week', f1) AS week,
date_part('dow', f1) AS dow,
date_part('isodow', f1) AS isodow,
date_part('doy', f1) AS doy,
-- date_part('julian', f1) AS julian,
date_part('epoch', f1) AS epoch
FROM date_tbl;

DROP TABLE DATE_TBL;

Expand All @@ -252,62 +252,62 @@ DROP TABLE DATE_TBL;
--@ SELECT EXTRACT(CENTURY FROM DATE '0101-12-31 BC'); -- -2
--@ SELECT EXTRACT(CENTURY FROM DATE '0100-12-31 BC'); -- -1
--@ SELECT EXTRACT(CENTURY FROM DATE '0001-12-31 BC'); -- -1
--@ SELECT EXTRACT(CENTURY FROM DATE '0001-01-01'); -- 1
SELECT EXTRACT(CENTURY FROM DATE '0001-01-01'); -- 1
--@ SELECT EXTRACT(CENTURY FROM DATE '0001-01-01 AD'); -- 1
--@ SELECT EXTRACT(CENTURY FROM DATE '1900-12-31'); -- 19
--@ SELECT EXTRACT(CENTURY FROM DATE '1901-01-01'); -- 20
--@ SELECT EXTRACT(CENTURY FROM DATE '2000-12-31'); -- 20
--@ SELECT EXTRACT(CENTURY FROM DATE '2001-01-01'); -- 21
SELECT EXTRACT(CENTURY FROM DATE '1900-12-31'); -- 19
SELECT EXTRACT(CENTURY FROM DATE '1901-01-01'); -- 20
SELECT EXTRACT(CENTURY FROM DATE '2000-12-31'); -- 20
SELECT EXTRACT(CENTURY FROM DATE '2001-01-01'); -- 21
--@ SELECT EXTRACT(CENTURY FROM CURRENT_DATE)>=21 AS True; -- true
--
-- millennium
--
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '0001-12-31 BC'); -- -1
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '0001-01-01 AD'); -- 1
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '1000-12-31'); -- 1
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '1001-01-01'); -- 2
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '2000-12-31'); -- 2
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '2001-01-01'); -- 3
--@ -- next test to be fixed on the turn of the next millennium;-)
SELECT EXTRACT(MILLENNIUM FROM DATE '1000-12-31'); -- 1
SELECT EXTRACT(MILLENNIUM FROM DATE '1001-01-01'); -- 2
SELECT EXTRACT(MILLENNIUM FROM DATE '2000-12-31'); -- 2
SELECT EXTRACT(MILLENNIUM FROM DATE '2001-01-01'); -- 3
-- next test to be fixed on the turn of the next millennium;-)
--@ SELECT EXTRACT(MILLENNIUM FROM CURRENT_DATE); -- 3
--
-- decade
--
--@ SELECT EXTRACT(DECADE FROM DATE '1994-12-25'); -- 199
--@ SELECT EXTRACT(DECADE FROM DATE '0010-01-01'); -- 1
--@ SELECT EXTRACT(DECADE FROM DATE '0009-12-31'); -- 0
SELECT EXTRACT(DECADE FROM DATE '1994-12-25'); -- 199
SELECT EXTRACT(DECADE FROM DATE '0010-01-01'); -- 1
SELECT EXTRACT(DECADE FROM DATE '0009-12-31'); -- 0
--@ SELECT EXTRACT(DECADE FROM DATE '0001-01-01 BC'); -- 0
--@ SELECT EXTRACT(DECADE FROM DATE '0002-12-31 BC'); -- -1
--@ SELECT EXTRACT(DECADE FROM DATE '0011-01-01 BC'); -- -1
--@ SELECT EXTRACT(DECADE FROM DATE '0012-12-31 BC'); -- -2
--
-- all possible fields
--
--@ SELECT EXTRACT(MICROSECONDS FROM DATE '2020-08-11');
--@ SELECT EXTRACT(MILLISECONDS FROM DATE '2020-08-11');
--@ SELECT EXTRACT(SECOND FROM DATE '2020-08-11');
--@ SELECT EXTRACT(MINUTE FROM DATE '2020-08-11');
--@ SELECT EXTRACT(HOUR FROM DATE '2020-08-11');
SELECT EXTRACT(MICROSECONDS FROM DATE '2020-08-11');
SELECT EXTRACT(MILLISECONDS FROM DATE '2020-08-11');
SELECT EXTRACT(SECOND FROM DATE '2020-08-11');
SELECT EXTRACT(MINUTE FROM DATE '2020-08-11');
SELECT EXTRACT(HOUR FROM DATE '2020-08-11');
SELECT EXTRACT(DAY FROM DATE '2020-08-11');
SELECT EXTRACT(MONTH FROM DATE '2020-08-11');
SELECT EXTRACT(YEAR FROM DATE '2020-08-11');
--@ SELECT EXTRACT(YEAR FROM DATE '2020-08-11 BC');
--@ SELECT EXTRACT(DECADE FROM DATE '2020-08-11');
--@ SELECT EXTRACT(CENTURY FROM DATE '2020-08-11');
--@ SELECT EXTRACT(MILLENNIUM FROM DATE '2020-08-11');
--@ SELECT EXTRACT(ISOYEAR FROM DATE '2020-08-11');
SELECT EXTRACT(DECADE FROM DATE '2020-08-11');
SELECT EXTRACT(CENTURY FROM DATE '2020-08-11');
SELECT EXTRACT(MILLENNIUM FROM DATE '2020-08-11');
SELECT EXTRACT(ISOYEAR FROM DATE '2020-08-11');
--@ SELECT EXTRACT(ISOYEAR FROM DATE '2020-08-11 BC');
--@ SELECT EXTRACT(QUARTER FROM DATE '2020-08-11');
--@ SELECT EXTRACT(WEEK FROM DATE '2020-08-11');
SELECT EXTRACT(QUARTER FROM DATE '2020-08-11');
SELECT EXTRACT(WEEK FROM DATE '2020-08-11');
SELECT EXTRACT(DOW FROM DATE '2020-08-11');
SELECT EXTRACT(DOW FROM DATE '2020-08-16');
--@ SELECT EXTRACT(ISODOW FROM DATE '2020-08-11');
--@ SELECT EXTRACT(ISODOW FROM DATE '2020-08-16');
SELECT EXTRACT(ISODOW FROM DATE '2020-08-11');
SELECT EXTRACT(ISODOW FROM DATE '2020-08-16');
SELECT EXTRACT(DOY FROM DATE '2020-08-11');
--@ SELECT EXTRACT(TIMEZONE FROM DATE '2020-08-11');
--@ SELECT EXTRACT(TIMEZONE_M FROM DATE '2020-08-11');
--@ SELECT EXTRACT(TIMEZONE_H FROM DATE '2020-08-11');
--@ SELECT EXTRACT(EPOCH FROM DATE '2020-08-11');
SELECT EXTRACT(EPOCH FROM DATE '2020-08-11');
--@ SELECT EXTRACT(JULIAN FROM DATE '2020-08-11');
--
-- test trunc function!
Expand Down
16 changes: 8 additions & 8 deletions src/tests/regress/data/sql/time.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,19 +61,19 @@ DROP TABLE TIME_TBL;
--
-- test EXTRACT
--
--@ SELECT EXTRACT(MICROSECOND FROM TIME '13:30:25.575401');
--@ SELECT EXTRACT(MILLISECOND FROM TIME '13:30:25.575401');
--@ SELECT EXTRACT(SECOND FROM TIME '13:30:25.575401');
SELECT EXTRACT(MICROSECOND FROM TIME '13:30:25.575401');
SELECT EXTRACT(MILLISECOND FROM TIME '13:30:25.575401');
SELECT EXTRACT(SECOND FROM TIME '13:30:25.575401');
SELECT EXTRACT(MINUTE FROM TIME '13:30:25.575401');
SELECT EXTRACT(HOUR FROM TIME '13:30:25.575401');
SELECT EXTRACT(DAY FROM TIME '13:30:25.575401'); -- error
SELECT EXTRACT(FORTNIGHT FROM TIME '13:30:25.575401'); -- error
SELECT EXTRACT(TIMEZONE FROM TIME '13:30:25.575401'); -- error
--@ SELECT EXTRACT(EPOCH FROM TIME '13:30:25.575401');
SELECT EXTRACT(EPOCH FROM TIME '13:30:25.575401');

-- date_part implementation is mostly the same as extract, so only
-- test a few cases for additional coverage.
--@ SELECT date_part('microsecond', TIME '13:30:25.575401');
--@ SELECT date_part('millisecond', TIME '13:30:25.575401');
--@ SELECT date_part('second', TIME '13:30:25.575401');
--@ SELECT date_part('epoch', TIME '13:30:25.575401');
SELECT date_part('microsecond', TIME '13:30:25.575401');
SELECT date_part('millisecond', TIME '13:30:25.575401');
SELECT date_part('second', TIME '13:30:25.575401');
SELECT date_part('epoch', TIME '13:30:25.575401');

0 comments on commit 795a19f

Please sign in to comment.