risingwave_expr/window_function/
call.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 std::fmt::Display;
16
17use FrameBound::{CurrentRow, Following, Preceding, UnboundedFollowing, UnboundedPreceding};
18use enum_as_inner::EnumAsInner;
19use parse_display::Display;
20use risingwave_common::types::DataType;
21use risingwave_common::util::sort_util::OrderType;
22use risingwave_common::{bail, must_match};
23use risingwave_pb::expr::window_frame::{PbBounds, PbExclusion};
24use risingwave_pb::expr::{PbWindowFrame, PbWindowFunction};
25
26use super::{
27    RangeFrameBounds, RowsFrameBound, RowsFrameBounds, SessionFrameBounds, WindowFuncKind,
28};
29use crate::Result;
30use crate::aggregate::AggArgs;
31
32#[derive(Debug, Clone)]
33pub struct WindowFuncCall {
34    pub kind: WindowFuncKind,
35    pub return_type: DataType,
36    pub args: AggArgs,
37    pub ignore_nulls: bool,
38    pub frame: Frame,
39}
40
41impl WindowFuncCall {
42    pub fn from_protobuf(call: &PbWindowFunction) -> Result<Self> {
43        let call = WindowFuncCall {
44            kind: WindowFuncKind::from_protobuf(call.get_type()?)?,
45            return_type: DataType::from(call.get_return_type()?),
46            args: AggArgs::from_protobuf(call.get_args())?,
47            ignore_nulls: call.get_ignore_nulls(),
48            frame: Frame::from_protobuf(call.get_frame()?)?,
49        };
50        Ok(call)
51    }
52}
53
54#[derive(Debug, Clone, Eq, PartialEq, Hash)]
55pub struct Frame {
56    pub bounds: FrameBounds,
57    pub exclusion: FrameExclusion,
58}
59
60impl Display for Frame {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        write!(f, "{}", self.bounds)?;
63        if self.exclusion != FrameExclusion::default() {
64            write!(f, " {}", self.exclusion)?;
65        }
66        Ok(())
67    }
68}
69
70impl Frame {
71    pub fn rows(start: RowsFrameBound, end: RowsFrameBound) -> Self {
72        Self {
73            bounds: FrameBounds::Rows(RowsFrameBounds { start, end }),
74            exclusion: FrameExclusion::default(),
75        }
76    }
77
78    pub fn rows_with_exclusion(
79        start: RowsFrameBound,
80        end: RowsFrameBound,
81        exclusion: FrameExclusion,
82    ) -> Self {
83        Self {
84            bounds: FrameBounds::Rows(RowsFrameBounds { start, end }),
85            exclusion,
86        }
87    }
88}
89
90impl Frame {
91    pub fn from_protobuf(frame: &PbWindowFrame) -> Result<Self> {
92        use risingwave_pb::expr::window_frame::PbType;
93        let bounds = match frame.get_type()? {
94            PbType::Unspecified => bail!("unspecified type of `WindowFrame`"),
95            #[expect(deprecated)]
96            PbType::RowsLegacy => {
97                #[expect(deprecated)]
98                {
99                    let start = FrameBound::<usize>::from_protobuf_legacy(frame.get_start()?)?;
100                    let end = FrameBound::<usize>::from_protobuf_legacy(frame.get_end()?)?;
101                    FrameBounds::Rows(RowsFrameBounds { start, end })
102                }
103            }
104            PbType::Rows => {
105                let bounds = must_match!(frame.get_bounds()?, PbBounds::Rows(bounds) => bounds);
106                FrameBounds::Rows(RowsFrameBounds::from_protobuf(bounds)?)
107            }
108            PbType::Range => {
109                let bounds = must_match!(frame.get_bounds()?, PbBounds::Range(bounds) => bounds);
110                FrameBounds::Range(RangeFrameBounds::from_protobuf(bounds)?)
111            }
112            PbType::Session => {
113                let bounds = must_match!(frame.get_bounds()?, PbBounds::Session(bounds) => bounds);
114                FrameBounds::Session(SessionFrameBounds::from_protobuf(bounds)?)
115            }
116        };
117        let exclusion = FrameExclusion::from_protobuf(frame.get_exclusion()?)?;
118        Ok(Self { bounds, exclusion })
119    }
120
121    pub fn to_protobuf(&self) -> PbWindowFrame {
122        use risingwave_pb::expr::window_frame::PbType;
123        let exclusion = self.exclusion.to_protobuf() as _;
124        #[expect(deprecated)] // because of `start` and `end` fields
125        match &self.bounds {
126            FrameBounds::Rows(bounds) => PbWindowFrame {
127                r#type: PbType::Rows as _,
128                start: None, // deprecated
129                end: None,   // deprecated
130                exclusion,
131                bounds: Some(PbBounds::Rows(bounds.to_protobuf())),
132            },
133            FrameBounds::Range(bounds) => PbWindowFrame {
134                r#type: PbType::Range as _,
135                start: None, // deprecated
136                end: None,   // deprecated
137                exclusion,
138                bounds: Some(PbBounds::Range(bounds.to_protobuf())),
139            },
140            FrameBounds::Session(bounds) => PbWindowFrame {
141                r#type: PbType::Session as _,
142                start: None, // deprecated
143                end: None,   // deprecated
144                exclusion,
145                bounds: Some(PbBounds::Session(bounds.to_protobuf())),
146            },
147        }
148    }
149}
150
151#[derive(Display, Debug, Clone, Eq, PartialEq, Hash, EnumAsInner)]
152#[display("{0}")]
153pub enum FrameBounds {
154    Rows(RowsFrameBounds),
155    // Groups(GroupsFrameBounds),
156    Range(RangeFrameBounds),
157    Session(SessionFrameBounds),
158}
159
160impl FrameBounds {
161    pub fn validate(&self) -> Result<()> {
162        match self {
163            Self::Rows(bounds) => bounds.validate(),
164            Self::Range(bounds) => bounds.validate(),
165            Self::Session(bounds) => bounds.validate(),
166        }
167    }
168
169    pub fn start_is_unbounded(&self) -> bool {
170        match self {
171            Self::Rows(RowsFrameBounds { start, .. }) => start.is_unbounded_preceding(),
172            Self::Range(RangeFrameBounds { start, .. }) => start.is_unbounded_preceding(),
173            Self::Session(_) => false,
174        }
175    }
176
177    pub fn end_is_unbounded(&self) -> bool {
178        match self {
179            Self::Rows(RowsFrameBounds { end, .. }) => end.is_unbounded_following(),
180            Self::Range(RangeFrameBounds { end, .. }) => end.is_unbounded_following(),
181            Self::Session(_) => false,
182        }
183    }
184
185    pub fn is_unbounded(&self) -> bool {
186        self.start_is_unbounded() || self.end_is_unbounded()
187    }
188
189    /// Whether the frame may include rows *preceding* the current row in the frame ordering.
190    pub fn may_include_preceding_rows(&self) -> bool {
191        match self {
192            Self::Rows(RowsFrameBounds { start, .. }) => {
193                start.is_unbounded_preceding() || start.is_preceding()
194            }
195            Self::Range(RangeFrameBounds { start, .. }) => {
196                start.is_unbounded_preceding() || start.is_preceding()
197            }
198            Self::Session(_) => true,
199        }
200    }
201
202    /// Whether the frame may include rows *following* the current row in the frame ordering.
203    pub fn may_include_following_rows(&self) -> bool {
204        match self {
205            Self::Rows(RowsFrameBounds { end, .. }) => {
206                end.is_unbounded_following() || end.is_following()
207            }
208            Self::Range(RangeFrameBounds { end, .. }) => {
209                end.is_unbounded_following() || end.is_following()
210            }
211            Self::Session(_) => true,
212        }
213    }
214}
215
216/// Whether a watermark on the first `ORDER BY` column can be forwarded through a general (i.e.
217/// non-EOWC) streaming over window operator evaluating the given `frames`, which all share the
218/// same `PARTITION BY` and `ORDER BY` clauses.
219///
220/// A watermark `wm` on a column promises that no row with the column value `< wm` will ever arrive
221/// again. To forward it, the operator must guarantee that the outputs of rows with the column value
222/// `< wm` will never change again, which holds when a row can only affect (i.e. be included in the
223/// frames of) rows that are not "smaller" than itself in the column:
224///
225/// - the column is ordered `ASC` and no frame extends to following rows, or
226/// - the column is ordered `DESC` and no frame extends to preceding rows.
227///
228/// Additionally, NULLs must be ordered as the largest values (`ASC NULLS LAST` or `DESC NULLS
229/// FIRST`). Rows with NULL in the column are not covered by the watermark guarantee, so they must
230/// also land on the "larger" side to never affect rows below the watermark.
231pub fn can_forward_watermark_on_order_key<'a>(
232    frames: impl IntoIterator<Item = &'a Frame>,
233    first_order_type: OrderType,
234) -> bool {
235    if !first_order_type.nulls_are_largest() {
236        return false;
237    }
238    let mut frames = frames.into_iter();
239    if first_order_type.is_ascending() {
240        frames.all(|frame| !frame.bounds.may_include_following_rows())
241    } else {
242        frames.all(|frame| !frame.bounds.may_include_preceding_rows())
243    }
244}
245
246pub trait FrameBoundsImpl {
247    fn validate(&self) -> Result<()>;
248}
249
250#[derive(Display, Debug, Clone, Eq, PartialEq, Hash, EnumAsInner)]
251#[display(style = "TITLE CASE")]
252pub enum FrameBound<T> {
253    UnboundedPreceding,
254    #[display("{0} PRECEDING")]
255    Preceding(T),
256    CurrentRow,
257    #[display("{0} FOLLOWING")]
258    Following(T),
259    UnboundedFollowing,
260}
261
262impl<T> FrameBound<T> {
263    fn offset_value(&self) -> Option<&T> {
264        match self {
265            UnboundedPreceding | UnboundedFollowing | CurrentRow => None,
266            Preceding(offset) | Following(offset) => Some(offset),
267        }
268    }
269
270    pub(super) fn validate_bounds(
271        start: &Self,
272        end: &Self,
273        offset_checker: impl Fn(&T) -> Result<()>,
274    ) -> Result<()> {
275        match (start, end) {
276            (_, UnboundedPreceding) => bail!("frame end cannot be UNBOUNDED PRECEDING"),
277            (UnboundedFollowing, _) => {
278                bail!("frame start cannot be UNBOUNDED FOLLOWING")
279            }
280            (Following(_), CurrentRow) | (Following(_), Preceding(_)) => {
281                bail!("frame starting from following row cannot have preceding rows")
282            }
283            (CurrentRow, Preceding(_)) => {
284                bail!("frame starting from current row cannot have preceding rows")
285            }
286            _ => {}
287        }
288
289        for bound in [start, end] {
290            if let Some(offset) = bound.offset_value() {
291                offset_checker(offset)?;
292            }
293        }
294
295        Ok(())
296    }
297
298    pub fn map<U>(self, f: impl Fn(T) -> U) -> FrameBound<U> {
299        match self {
300            UnboundedPreceding => UnboundedPreceding,
301            Preceding(offset) => Preceding(f(offset)),
302            CurrentRow => CurrentRow,
303            Following(offset) => Following(f(offset)),
304            UnboundedFollowing => UnboundedFollowing,
305        }
306    }
307}
308
309impl<T> FrameBound<T>
310where
311    T: Copy,
312{
313    pub(super) fn reverse(self) -> FrameBound<T> {
314        match self {
315            UnboundedPreceding => UnboundedFollowing,
316            Preceding(offset) => Following(offset),
317            CurrentRow => CurrentRow,
318            Following(offset) => Preceding(offset),
319            UnboundedFollowing => UnboundedPreceding,
320        }
321    }
322}
323
324#[derive(Display, Debug, Copy, Clone, Eq, PartialEq, Hash, Default, EnumAsInner)]
325#[display("EXCLUDE {}", style = "TITLE CASE")]
326pub enum FrameExclusion {
327    CurrentRow,
328    // Group,
329    // Ties,
330    #[default]
331    NoOthers,
332}
333
334impl FrameExclusion {
335    fn from_protobuf(exclusion: PbExclusion) -> Result<Self> {
336        let excl = match exclusion {
337            PbExclusion::Unspecified => bail!("unspecified type of `FrameExclusion`"),
338            PbExclusion::CurrentRow => Self::CurrentRow,
339            PbExclusion::NoOthers => Self::NoOthers,
340        };
341        Ok(excl)
342    }
343
344    fn to_protobuf(self) -> PbExclusion {
345        match self {
346            Self::CurrentRow => PbExclusion::CurrentRow,
347            Self::NoOthers => PbExclusion::NoOthers,
348        }
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use risingwave_common::types::ScalarImpl;
355
356    use super::*;
357    use crate::window_function::{RangeFrameBound, RangeFrameOffset, SessionFrameGap};
358
359    fn rows_frame(start: RowsFrameBound, end: RowsFrameBound) -> Frame {
360        Frame::rows(start, end)
361    }
362
363    fn range_frame(start: FrameBound<i64>, end: FrameBound<i64>, order_type: OrderType) -> Frame {
364        // The offsets don't need to be prepared for evaluation here.
365        let bound = |bound: FrameBound<i64>| -> RangeFrameBound {
366            bound.map(|offset| RangeFrameOffset::new(ScalarImpl::Int64(offset)))
367        };
368        Frame {
369            bounds: FrameBounds::Range(RangeFrameBounds {
370                order_data_type: DataType::Int64,
371                order_type,
372                offset_data_type: DataType::Int64,
373                start: bound(start),
374                end: bound(end),
375            }),
376            exclusion: FrameExclusion::default(),
377        }
378    }
379
380    fn session_frame(order_type: OrderType) -> Frame {
381        Frame {
382            bounds: FrameBounds::Session(SessionFrameBounds {
383                order_data_type: DataType::Int64,
384                order_type,
385                gap_data_type: DataType::Int64,
386                gap: SessionFrameGap::new(ScalarImpl::Int64(10)),
387            }),
388            exclusion: FrameExclusion::default(),
389        }
390    }
391
392    #[test]
393    fn test_can_forward_watermark_on_order_key() {
394        let asc = OrderType::ascending(); // ASC NULLS LAST
395        let desc = OrderType::descending(); // DESC NULLS FIRST
396        let can_forward =
397            |frames: &[Frame], order_type| can_forward_watermark_on_order_key(frames, order_type);
398
399        let preceding_only = [
400            rows_frame(UnboundedPreceding, CurrentRow), // rank functions
401            rows_frame(Preceding(1), Preceding(1)),     // `lag`
402            rows_frame(Preceding(2), CurrentRow),
403        ];
404        let following_only = [
405            rows_frame(CurrentRow, UnboundedFollowing),
406            rows_frame(Following(1), Following(1)), // `lead`
407            rows_frame(CurrentRow, Following(2)),
408        ];
409        let current_row_only = [rows_frame(CurrentRow, CurrentRow)];
410        let both_sides = [
411            rows_frame(Preceding(1), Following(1)),
412            rows_frame(UnboundedPreceding, UnboundedFollowing),
413        ];
414
415        assert!(can_forward(&preceding_only, asc));
416        assert!(!can_forward(&preceding_only, desc));
417        assert!(!can_forward(&following_only, asc));
418        assert!(can_forward(&following_only, desc));
419        assert!(can_forward(&current_row_only, asc));
420        assert!(can_forward(&current_row_only, desc));
421        assert!(!can_forward(&both_sides, asc));
422        assert!(!can_forward(&both_sides, desc));
423
424        // Any frame extending to the "smaller" side of the current row disables forwarding.
425        let mixed = [
426            rows_frame(Preceding(1), Preceding(1)),
427            rows_frame(Following(1), Following(1)),
428        ];
429        assert!(!can_forward(&mixed, asc));
430        assert!(!can_forward(&mixed, desc));
431
432        // `RANGE` frames follow the same rule.
433        for order_type in [asc, desc] {
434            let preceding_only = [range_frame(Preceding(10), CurrentRow, order_type)];
435            let following_only = [range_frame(CurrentRow, Following(10), order_type)];
436            let both_sides = [range_frame(Preceding(10), Following(10), order_type)];
437            assert_eq!(
438                can_forward(&preceding_only, order_type),
439                order_type.is_ascending()
440            );
441            assert_eq!(
442                can_forward(&following_only, order_type),
443                order_type.is_descending()
444            );
445            assert!(!can_forward(&both_sides, order_type));
446        }
447
448        // `SESSION` frames extend to both sides.
449        for order_type in [asc, desc] {
450            assert!(!can_forward(&[session_frame(order_type)], order_type));
451        }
452
453        // NULLs must be ordered as the largest values.
454        assert!(!can_forward(
455            &preceding_only,
456            OrderType::ascending_nulls_first()
457        ));
458        assert!(!can_forward(
459            &following_only,
460            OrderType::descending_nulls_last()
461        ));
462        assert!(can_forward(
463            &preceding_only,
464            OrderType::ascending_nulls_last()
465        ));
466        assert!(can_forward(
467            &following_only,
468            OrderType::descending_nulls_first()
469        ));
470    }
471}
risingwave_expr/window_function/call.rs

risingwave_expr/window_function/
call.rs
