risingwave_expr/window_function/
range.rs

1// Copyright 2024 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;
16use std::ops::Deref;
17use std::sync::Arc;
18
19use anyhow::Context;
20use educe::Educe;
21use risingwave_common::bail;
22use risingwave_common::row::OwnedRow;
23use risingwave_common::types::{
24    DataType, Datum, IsNegative, ScalarImpl, ScalarRefImpl, Sentinelled, ToOwnedDatum, ToText,
25};
26use risingwave_common::util::sort_util::{Direction, OrderType};
27use risingwave_common::util::value_encoding::{DatumFromProtoExt, DatumToProtoExt};
28use risingwave_pb::expr::window_frame::{PbBoundType, PbRangeFrameBound, PbRangeFrameBounds};
29
30use super::FrameBound::{
31    self, CurrentRow, Following, Preceding, UnboundedFollowing, UnboundedPreceding,
32};
33use super::FrameBoundsImpl;
34use crate::Result;
35use crate::expr::{
36    InputRefExpression, LiteralExpression, SyncExpression, SyncExpressionBoxExt, build_func,
37};
38
39#[derive(Debug, Clone, Eq, PartialEq, Hash)]
40pub struct RangeFrameBounds {
41    pub order_data_type: DataType,
42    pub order_type: OrderType,
43    pub offset_data_type: DataType,
44    pub start: RangeFrameBound,
45    pub end: RangeFrameBound,
46}
47
48impl RangeFrameBounds {
49    pub(super) fn from_protobuf(bounds: &PbRangeFrameBounds) -> Result<Self> {
50        let order_data_type = DataType::from(bounds.get_order_data_type()?);
51        let order_type = OrderType::from_protobuf(bounds.get_order_type()?);
52        let offset_data_type = DataType::from(bounds.get_offset_data_type()?);
53        let start = FrameBound::<RangeFrameOffset>::from_protobuf(
54            bounds.get_start()?,
55            &order_data_type,
56            &offset_data_type,
57        )?;
58        let end = FrameBound::<RangeFrameOffset>::from_protobuf(
59            bounds.get_end()?,
60            &order_data_type,
61            &offset_data_type,
62        )?;
63        Ok(Self {
64            order_data_type,
65            order_type,
66            offset_data_type,
67            start,
68            end,
69        })
70    }
71
72    pub(super) fn to_protobuf(&self) -> PbRangeFrameBounds {
73        PbRangeFrameBounds {
74            start: Some(self.start.to_protobuf()),
75            end: Some(self.end.to_protobuf()),
76            order_data_type: Some(self.order_data_type.to_protobuf()),
77            order_type: Some(self.order_type.to_protobuf()),
78            offset_data_type: Some(self.offset_data_type.to_protobuf()),
79        }
80    }
81}
82
83impl Display for RangeFrameBounds {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        write!(
86            f,
87            "RANGE BETWEEN {} AND {}",
88            self.start.for_display(),
89            self.end.for_display()
90        )?;
91        Ok(())
92    }
93}
94
95impl FrameBoundsImpl for RangeFrameBounds {
96    fn validate(&self) -> Result<()> {
97        fn validate_non_negative(val: impl IsNegative + Display) -> Result<()> {
98            if val.is_negative() {
99                bail!(
100                    "frame bound offset should be non-negative, but {} is given",
101                    val
102                );
103            }
104            Ok(())
105        }
106
107        FrameBound::validate_bounds(&self.start, &self.end, |offset| {
108            match offset.as_scalar_ref_impl() {
109                // TODO(rc): use decl macro?
110                ScalarRefImpl::Int16(val) => validate_non_negative(val)?,
111                ScalarRefImpl::Int32(val) => validate_non_negative(val)?,
112                ScalarRefImpl::Int64(val) => validate_non_negative(val)?,
113                ScalarRefImpl::Float32(val) => validate_non_negative(val)?,
114                ScalarRefImpl::Float64(val) => validate_non_negative(val)?,
115                ScalarRefImpl::Decimal(val) => validate_non_negative(val)?,
116                ScalarRefImpl::Interval(val) => {
117                    if !val.is_never_negative() {
118                        bail!(
119                            "for frame bound offset of type `interval`, each field should be non-negative, but {} is given",
120                            val
121                        );
122                    }
123                    if matches!(self.order_data_type, DataType::Timestamptz) {
124                        // for `timestamptz`, we only support offset without `month` and `day` fields
125                        if val.months() != 0 || val.days() != 0 {
126                            bail!(
127                                "for frame order column of type `timestamptz`, offset should not have non-zero `month` and `day`",
128                            );
129                        }
130                    }
131                }
132                _ => unreachable!(
133                    "other order column data types are not supported and should be banned in frontend"
134                ),
135            }
136            Ok(())
137        })
138    }
139}
140
141impl RangeFrameBounds {
142    /// Get the frame start for a given order column value.
143    ///
144    /// ## Examples
145    ///
146    /// For the following frames:
147    ///
148    /// ```sql
149    /// ORDER BY x ASC RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
150    /// ORDER BY x DESC RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
151    /// ```
152    ///
153    /// For any CURRENT ROW with any order value, the frame start is always the first-most row, which is
154    /// represented by [`Sentinelled::Smallest`].
155    ///
156    /// For the following frame:
157    ///
158    /// ```sql
159    /// ORDER BY x ASC RANGE BETWEEN 10 PRECEDING AND CURRENT ROW
160    /// ```
161    ///
162    /// For CURRENT ROW with order value `100`, the frame start is the **FIRST** row with order value `90`.
163    ///
164    /// For the following frame:
165    ///
166    /// ```sql
167    /// ORDER BY x DESC RANGE BETWEEN 10 PRECEDING AND CURRENT ROW
168    /// ```
169    ///
170    /// For CURRENT ROW with order value `100`, the frame start is the **FIRST** row with order value `110`.
171    pub fn frame_start_of(&self, order_value: impl ToOwnedDatum) -> Sentinelled<Datum> {
172        self.start.for_calc().bound_of(order_value, self.order_type)
173    }
174
175    /// Get the frame end for a given order column value. It's very similar to `frame_start_of`, just with
176    /// everything on the other direction.
177    pub fn frame_end_of(&self, order_value: impl ToOwnedDatum) -> Sentinelled<Datum> {
178        self.end.for_calc().bound_of(order_value, self.order_type)
179    }
180
181    /// Get the order value of the CURRENT ROW of the first frame that includes the given order value.
182    ///
183    /// ## Examples
184    ///
185    /// For the following frames:
186    ///
187    /// ```sql
188    /// ORDER BY x ASC RANGE BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING
189    /// ORDER BY x DESC RANGE BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING
190    /// ```
191    ///
192    /// For any given order value, the first CURRENT ROW is always the first-most row, which is
193    /// represented by [`Sentinelled::Smallest`].
194    ///
195    /// For the following frame:
196    ///
197    /// ```sql
198    /// ORDER BY x ASC RANGE BETWEEN CURRENT ROW AND 10 FOLLOWING
199    /// ```
200    ///
201    /// For a given order value `100`, the first CURRENT ROW should have order value `90`.
202    ///
203    /// For the following frame:
204    ///
205    /// ```sql
206    /// ORDER BY x DESC RANGE BETWEEN CURRENT ROW AND 10 FOLLOWING
207    /// ```
208    ///
209    /// For a given order value `100`, the first CURRENT ROW should have order value `110`.
210    pub fn first_curr_of(&self, order_value: impl ToOwnedDatum) -> Sentinelled<Datum> {
211        self.end
212            .for_calc()
213            .reverse()
214            .bound_of(order_value, self.order_type)
215    }
216
217    /// Get the order value of the CURRENT ROW of the last frame that includes the given order value.
218    /// It's very similar to `first_curr_of`, just with everything on the other direction.
219    pub fn last_curr_of(&self, order_value: impl ToOwnedDatum) -> Sentinelled<Datum> {
220        self.start
221            .for_calc()
222            .reverse()
223            .bound_of(order_value, self.order_type)
224    }
225}
226
227pub type RangeFrameBound = FrameBound<RangeFrameOffset>;
228
229impl RangeFrameBound {
230    fn from_protobuf(
231        bound: &PbRangeFrameBound,
232        order_data_type: &DataType,
233        offset_data_type: &DataType,
234    ) -> Result<Self> {
235        let bound = match bound.get_type()? {
236            PbBoundType::Unspecified => bail!("unspecified type of `RangeFrameBound`"),
237            PbBoundType::UnboundedPreceding => Self::UnboundedPreceding,
238            PbBoundType::CurrentRow => Self::CurrentRow,
239            PbBoundType::UnboundedFollowing => Self::UnboundedFollowing,
240            bound_type @ (PbBoundType::Preceding | PbBoundType::Following) => {
241                let offset_value = Datum::from_protobuf(bound.get_offset()?, offset_data_type)
242                    .context("offset `Datum` is not decodable")?
243                    .context("offset of `RangeFrameBound` must be non-NULL")?;
244                let mut offset = RangeFrameOffset::new(offset_value);
245                offset.prepare(order_data_type, offset_data_type)?;
246                if bound_type == PbBoundType::Preceding {
247                    Self::Preceding(offset)
248                } else {
249                    Self::Following(offset)
250                }
251            }
252        };
253        Ok(bound)
254    }
255
256    fn to_protobuf(&self) -> PbRangeFrameBound {
257        let (r#type, offset) = match self {
258            Self::UnboundedPreceding => (PbBoundType::UnboundedPreceding, None),
259            Self::Preceding(offset) => (
260                PbBoundType::Preceding,
261                Some(Some(offset.as_scalar_ref_impl()).to_protobuf()),
262            ),
263            Self::CurrentRow => (PbBoundType::CurrentRow, None),
264            Self::Following(offset) => (
265                PbBoundType::Following,
266                Some(Some(offset.as_scalar_ref_impl()).to_protobuf()),
267            ),
268            Self::UnboundedFollowing => (PbBoundType::UnboundedFollowing, None),
269        };
270        PbRangeFrameBound {
271            r#type: r#type as _,
272            offset,
273        }
274    }
275}
276
277impl RangeFrameBound {
278    fn for_display(&self) -> FrameBound<String> {
279        match self {
280            UnboundedPreceding => UnboundedPreceding,
281            Preceding(offset) => Preceding(offset.as_scalar_ref_impl().to_text()),
282            CurrentRow => CurrentRow,
283            Following(offset) => Following(offset.as_scalar_ref_impl().to_text()),
284            UnboundedFollowing => UnboundedFollowing,
285        }
286    }
287
288    fn for_calc(&self) -> FrameBound<RangeFrameOffsetRef<'_>> {
289        match self {
290            UnboundedPreceding => UnboundedPreceding,
291            Preceding(offset) => Preceding(RangeFrameOffsetRef {
292                add_expr: offset.add_expr.as_ref().unwrap().as_ref(),
293                sub_expr: offset.sub_expr.as_ref().unwrap().as_ref(),
294            }),
295            CurrentRow => CurrentRow,
296            Following(offset) => Following(RangeFrameOffsetRef {
297                add_expr: offset.add_expr.as_ref().unwrap().as_ref(),
298                sub_expr: offset.sub_expr.as_ref().unwrap().as_ref(),
299            }),
300            UnboundedFollowing => UnboundedFollowing,
301        }
302    }
303}
304
305/// The wrapper type for [`ScalarImpl`] range frame offset, containing
306/// two expressions to help adding and subtracting the offset.
307#[derive(Debug, Clone, Educe)]
308#[educe(PartialEq, Eq, Hash)]
309pub struct RangeFrameOffset {
310    /// The original offset value.
311    offset: ScalarImpl,
312    /// Built expression for `$0 + offset`.
313    #[educe(PartialEq(ignore), Hash(ignore))]
314    add_expr: Option<Arc<dyn SyncExpression>>,
315    /// Built expression for `$0 - offset`.
316    #[educe(PartialEq(ignore), Hash(ignore))]
317    sub_expr: Option<Arc<dyn SyncExpression>>,
318}
319
320impl RangeFrameOffset {
321    pub fn new(offset: ScalarImpl) -> Self {
322        Self {
323            offset,
324            add_expr: None,
325            sub_expr: None,
326        }
327    }
328
329    fn prepare(&mut self, order_data_type: &DataType, offset_data_type: &DataType) -> Result<()> {
330        use risingwave_pb::expr::expr_node::PbType as PbExprType;
331
332        let input_expr = InputRefExpression::new(order_data_type.clone(), 0);
333        let offset_expr =
334            LiteralExpression::new(offset_data_type.clone(), Some(self.offset.clone()));
335        let add_expr = build_func(
336            PbExprType::Add,
337            order_data_type.clone(),
338            vec![input_expr.clone().boxed(), offset_expr.clone().boxed()],
339        )?;
340        let crate::expr::BoxedExpression::Sync(add_expr) = add_expr else {
341            bail!("range frame offset add expression must be sync");
342        };
343        self.add_expr = Some(add_expr);
344
345        let sub_expr = build_func(
346            PbExprType::Subtract,
347            order_data_type.clone(),
348            vec![input_expr.boxed(), offset_expr.boxed()],
349        )?;
350        let crate::expr::BoxedExpression::Sync(sub_expr) = sub_expr else {
351            bail!("range frame offset subtract expression must be sync");
352        };
353        self.sub_expr = Some(sub_expr);
354        Ok(())
355    }
356
357    pub fn new_for_test(
358        offset: ScalarImpl,
359        order_data_type: &DataType,
360        offset_data_type: &DataType,
361    ) -> Self {
362        let mut offset = Self::new(offset);
363        offset.prepare(order_data_type, offset_data_type).unwrap();
364        offset
365    }
366}
367
368impl Deref for RangeFrameOffset {
369    type Target = ScalarImpl;
370
371    fn deref(&self) -> &Self::Target {
372        &self.offset
373    }
374}
375
376#[derive(Debug, Educe)]
377#[educe(Clone, Copy)]
378struct RangeFrameOffsetRef<'a> {
379    /// Built expression for `$0 + offset`.
380    add_expr: &'a dyn SyncExpression,
381    /// Built expression for `$0 - offset`.
382    sub_expr: &'a dyn SyncExpression,
383}
384
385impl FrameBound<RangeFrameOffsetRef<'_>> {
386    fn bound_of(self, order_value: impl ToOwnedDatum, order_type: OrderType) -> Sentinelled<Datum> {
387        let expr = match (self, order_type.direction()) {
388            (UnboundedPreceding, _) => return Sentinelled::Smallest,
389            (UnboundedFollowing, _) => return Sentinelled::Largest,
390            (CurrentRow, _) => return Sentinelled::Normal(order_value.to_owned_datum()),
391            (Preceding(offset), Direction::Ascending)
392            | (Following(offset), Direction::Descending) => {
393                // should SUBTRACT the offset
394                offset.sub_expr
395            }
396            (Following(offset), Direction::Ascending)
397            | (Preceding(offset), Direction::Descending) => {
398                // should ADD the offset
399                offset.add_expr
400            }
401        };
402        let row = OwnedRow::new(vec![order_value.to_owned_datum()]);
403        Sentinelled::Normal(
404            expr.eval_row(&row)
405                .expect("just simple calculation, should succeed"), // TODO(rc): handle overflow
406        )
407    }
408}
risingwave_expr/window_function/range.rs

risingwave_expr/window_function/
range.rs
