risingwave_expr_impl/scalar/
tumble.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use num_traits::Zero;
16use risingwave_common::types::{Date, Interval, Timestamp, Timestamptz};
17use risingwave_expr::{ExprError, Result, function};
18
19#[inline(always)]
20fn interval_to_micro_second(t: Interval) -> Result<i64> {
21    let checked_interval_to_micro_second = || {
22        (t.months() as i64)
23            .checked_mul(Interval::USECS_PER_MONTH)?
24            .checked_add(
25                (t.days() as i64)
26                    .checked_mul(Interval::USECS_PER_DAY)?
27                    .checked_add(t.usecs())?,
28            )
29    };
30
31    checked_interval_to_micro_second().ok_or(ExprError::NumericOutOfRange)
32}
33
34#[function("tumble_start(date, interval) -> timestamp")]
35pub fn tumble_start_date(timestamp: Date, window_size: Interval) -> Result<Timestamp> {
36    tumble_start_date_time(timestamp.into(), window_size)
37}
38
39#[function("tumble_start(timestamp, interval) -> timestamp")]
40pub fn tumble_start_date_time(timestamp: Timestamp, window_size: Interval) -> Result<Timestamp> {
41    let timestamp_micro_second = timestamp.0.and_utc().timestamp_micros();
42    let window_start_micro_second = get_window_start(timestamp_micro_second, window_size)?;
43    Timestamp::with_micros(window_start_micro_second).map_err(|_| ExprError::NumericOutOfRange)
44}
45
46#[function("tumble_start(timestamptz, interval) -> timestamptz")]
47pub fn tumble_start_timestamptz(tz: Timestamptz, window_size: Interval) -> Result<Timestamptz> {
48    get_window_start(tz.timestamp_micros(), window_size)
49        .and_then(|us| Timestamptz::from_micros(us).ok_or(ExprError::NumericOutOfRange))
50}
51
52/// The common part of PostgreSQL function `timestamp_bin` and `timestamptz_bin`.
53#[inline(always)]
54fn get_window_start(timestamp_micro_second: i64, window_size: Interval) -> Result<i64> {
55    get_window_start_with_offset(timestamp_micro_second, window_size, Interval::zero())
56}
57
58#[function("tumble_start(date, interval, interval) -> timestamp")]
59pub fn tumble_start_offset_date(
60    timestamp_date: Date,
61    window_size: Interval,
62    offset: Interval,
63) -> Result<Timestamp> {
64    tumble_start_offset_date_time(timestamp_date.into(), window_size, offset)
65}
66
67#[function("tumble_start(timestamp, interval, interval) -> timestamp")]
68pub fn tumble_start_offset_date_time(
69    time: Timestamp,
70    window_size: Interval,
71    offset: Interval,
72) -> Result<Timestamp> {
73    let timestamp_micro_second = time.0.and_utc().timestamp_micros();
74    let window_start_micro_second =
75        get_window_start_with_offset(timestamp_micro_second, window_size, offset)?;
76
77    Timestamp::with_micros(window_start_micro_second).map_err(|_| ExprError::NumericOutOfRange)
78}
79
80#[inline(always)]
81fn get_window_start_with_offset(
82    timestamp_micro_second: i64,
83    window_size: Interval,
84    offset: Interval,
85) -> Result<i64> {
86    let window_size_micro_second = interval_to_micro_second(window_size)?;
87    let offset_micro_second = interval_to_micro_second(offset)?;
88
89    // Inspired by https://issues.apache.org/jira/browse/FLINK-26334
90    let remainder = timestamp_micro_second
91        .checked_sub(offset_micro_second)
92        .ok_or(ExprError::NumericOutOfRange)?
93        .checked_rem(window_size_micro_second)
94        .ok_or(ExprError::DivisionByZero)?;
95    if remainder < 0 {
96        timestamp_micro_second
97            .checked_sub(remainder + window_size_micro_second)
98            .ok_or(ExprError::NumericOutOfRange)
99    } else {
100        timestamp_micro_second
101            .checked_sub(remainder)
102            .ok_or(ExprError::NumericOutOfRange)
103    }
104}
105
106#[function("tumble_start(timestamptz, interval, interval) -> timestamptz")]
107pub fn tumble_start_offset_timestamptz(
108    tz: Timestamptz,
109    window_size: Interval,
110    offset: Interval,
111) -> Result<Timestamptz> {
112    get_window_start_with_offset(tz.timestamp_micros(), window_size, offset)
113        .and_then(|us| Timestamptz::from_micros(us).ok_or(ExprError::NumericOutOfRange))
114}
115
116#[cfg(test)]
117mod tests {
118    use chrono::{Datelike, Timelike};
119    use risingwave_common::types::test_utils::IntervalTestExt;
120    use risingwave_common::types::{Date, Interval, Timestamp};
121
122    use super::tumble_start_offset_date_time;
123    use crate::scalar::tumble::{
124        get_window_start, interval_to_micro_second, tumble_start_date_time,
125    };
126
127    #[test]
128    fn test_tumble_start_date_time() {
129        let dt = Date::from_ymd_uncheck(2022, 2, 22).and_hms_uncheck(22, 22, 22);
130        let interval = Interval::from_minutes(30);
131        let w = tumble_start_date_time(dt, interval).unwrap().0;
132        assert_eq!(w.year(), 2022);
133        assert_eq!(w.month(), 2);
134        assert_eq!(w.day(), 22);
135        assert_eq!(w.hour(), 22);
136        assert_eq!(w.minute(), 0);
137        assert_eq!(w.second(), 0);
138    }
139
140    #[test]
141    fn test_tumble_start_negative_fractional() {
142        // A pre-1970 sub-second window start used to wrap `as u32` and panic.
143        let dt = "1969-12-31 23:59:58.450".parse::<Timestamp>().unwrap();
144        let w = tumble_start_date_time(dt, Interval::from_millis(100)).unwrap();
145        assert_eq!(w, "1969-12-31 23:59:58.400".parse::<Timestamp>().unwrap());
146    }
147
148    #[test]
149    fn test_tumble_start_offset_date_time() {
150        let dt = Date::from_ymd_uncheck(2022, 2, 22).and_hms_uncheck(22, 22, 22);
151        let window_size = 30;
152        for offset in 0..window_size {
153            for coefficient in 0..5 {
154                let w = tumble_start_date_time(dt, Interval::from_minutes(window_size))
155                    .unwrap()
156                    .0;
157                println!("{}", w);
158                let w = tumble_start_offset_date_time(
159                    dt,
160                    Interval::from_minutes(window_size),
161                    Interval::from_minutes(coefficient * window_size + offset),
162                )
163                .unwrap()
164                .0;
165                assert_eq!(w.year(), 2022);
166                assert_eq!(w.month(), 2);
167                assert_eq!(w.day(), 22);
168                if offset > 22 {
169                    assert_eq!(w.hour(), 21);
170                    assert_eq!(w.minute(), 30 + offset as u32);
171                } else {
172                    assert_eq!(w.hour(), 22);
173                    assert_eq!(w.minute(), offset as u32);
174                }
175
176                assert_eq!(w.second(), 0);
177            }
178        }
179    }
180
181    #[test]
182    fn test_remainder_necessary() {
183        let mut wrong_cnt = 0;
184        for i in -30..30 {
185            let timestamp_micro_second = Interval::from_minutes(i).usecs();
186            let window_size = Interval::from_minutes(5);
187            let window_start = get_window_start(timestamp_micro_second, window_size).unwrap();
188
189            let window_size_micro_second = interval_to_micro_second(window_size).unwrap();
190            let default_window_start = timestamp_micro_second
191                - (timestamp_micro_second + window_size_micro_second) % window_size_micro_second;
192
193            if timestamp_micro_second < default_window_start {
194                // which is wrong
195                wrong_cnt += 1;
196            }
197
198            assert!(timestamp_micro_second >= window_start)
199        }
200        assert_ne!(wrong_cnt, 0);
201    }
202
203    #[test]
204    fn test_window_start_overflow() {
205        get_window_start(i64::MIN, Interval::from_millis(20)).unwrap_err();
206        interval_to_micro_second(Interval::from_month_day_usec(1, 1, i64::MAX)).unwrap_err();
207    }
208}
risingwave_expr_impl/scalar/tumble.rs

risingwave_expr_impl/scalar/
tumble.rs
