risingwave_expr/window_function/
session.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, ToOwnedDatum, ToText,
25};
26use risingwave_common::util::sort_util::OrderType;
27use risingwave_common::util::value_encoding::{DatumFromProtoExt, DatumToProtoExt};
28use risingwave_pb::expr::window_frame::PbSessionFrameBounds;
29
30use super::FrameBoundsImpl;
31use crate::Result;
32use crate::expr::{
33    InputRefExpression, LiteralExpression, SyncExpression, SyncExpressionBoxExt, build_func,
34};
35
36/// To implement Session Window in a similar way to Range Frame, we define a similar frame bounds
37/// structure here. It's very like [`RangeFrameBounds`](super::RangeFrameBounds), but with a gap
38/// instead of start & end offset.
39#[derive(Debug, Clone, Eq, PartialEq, Hash)]
40pub struct SessionFrameBounds {
41    pub order_data_type: DataType,
42    pub order_type: OrderType,
43    pub gap_data_type: DataType,
44    pub gap: SessionFrameGap,
45}
46
47impl SessionFrameBounds {
48    pub(super) fn from_protobuf(bounds: &PbSessionFrameBounds) -> Result<Self> {
49        let order_data_type = DataType::from(bounds.get_order_data_type()?);
50        let order_type = OrderType::from_protobuf(bounds.get_order_type()?);
51        let gap_data_type = DataType::from(bounds.get_gap_data_type()?);
52        let gap_value = Datum::from_protobuf(bounds.get_gap()?, &gap_data_type)
53            .context("gap `Datum` is not decodable")?
54            .context("gap of session frame must be non-NULL")?;
55        let mut gap = SessionFrameGap::new(gap_value);
56        gap.prepare(&order_data_type, &gap_data_type)?;
57        Ok(Self {
58            order_data_type,
59            order_type,
60            gap_data_type,
61            gap,
62        })
63    }
64
65    pub(super) fn to_protobuf(&self) -> PbSessionFrameBounds {
66        PbSessionFrameBounds {
67            gap: Some(Some(self.gap.as_scalar_ref_impl()).to_protobuf()),
68            order_data_type: Some(self.order_data_type.to_protobuf()),
69            order_type: Some(self.order_type.to_protobuf()),
70            gap_data_type: Some(self.gap_data_type.to_protobuf()),
71        }
72    }
73}
74
75impl Display for SessionFrameBounds {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        write!(
78            f,
79            "SESSION WITH GAP {}",
80            self.gap.as_scalar_ref_impl().to_text()
81        )
82    }
83}
84
85impl FrameBoundsImpl for SessionFrameBounds {
86    fn validate(&self) -> Result<()> {
87        // TODO(rc): maybe can merge with `RangeFrameBounds::validate`
88
89        fn validate_non_negative(val: impl IsNegative + Display) -> Result<()> {
90            if val.is_negative() {
91                bail!("session gap should be non-negative, but {} is given", val);
92            }
93            Ok(())
94        }
95
96        match self.gap.as_scalar_ref_impl() {
97            ScalarRefImpl::Int16(val) => validate_non_negative(val)?,
98            ScalarRefImpl::Int32(val) => validate_non_negative(val)?,
99            ScalarRefImpl::Int64(val) => validate_non_negative(val)?,
100            ScalarRefImpl::Float32(val) => validate_non_negative(val)?,
101            ScalarRefImpl::Float64(val) => validate_non_negative(val)?,
102            ScalarRefImpl::Decimal(val) => validate_non_negative(val)?,
103            ScalarRefImpl::Interval(val) => {
104                if !val.is_never_negative() {
105                    bail!(
106                        "for session gap of type `interval`, each field should be non-negative, but {} is given",
107                        val
108                    );
109                }
110                if matches!(self.order_data_type, DataType::Timestamptz) {
111                    // for `timestamptz`, we only support gap without `month` and `day` fields
112                    if val.months() != 0 || val.days() != 0 {
113                        bail!(
114                            "for session order column of type `timestamptz`, gap should not have non-zero `month` and `day`",
115                        );
116                    }
117                }
118            }
119            _ => unreachable!(
120                "other order column data types are not supported and should be banned in frontend"
121            ),
122        }
123        Ok(())
124    }
125}
126
127impl SessionFrameBounds {
128    pub fn minimal_next_start_of(&self, end_order_value: impl ToOwnedDatum) -> Datum {
129        self.gap.for_calc().minimal_next_start_of(end_order_value)
130    }
131}
132
133/// The wrapper type for [`ScalarImpl`] session gap, containing an expression to help adding the gap
134/// to a given value.
135#[derive(Debug, Clone, Educe)]
136#[educe(PartialEq, Eq, Hash)]
137pub struct SessionFrameGap {
138    /// The original gap value.
139    gap: ScalarImpl,
140    /// Built expression for `$0 + gap`.
141    #[educe(PartialEq(ignore), Hash(ignore))]
142    add_expr: Option<Arc<dyn SyncExpression>>,
143}
144
145impl Deref for SessionFrameGap {
146    type Target = ScalarImpl;
147
148    fn deref(&self) -> &Self::Target {
149        &self.gap
150    }
151}
152
153impl SessionFrameGap {
154    pub fn new(gap: ScalarImpl) -> Self {
155        Self {
156            gap,
157            add_expr: None,
158        }
159    }
160
161    fn prepare(&mut self, order_data_type: &DataType, gap_data_type: &DataType) -> Result<()> {
162        use risingwave_pb::expr::expr_node::PbType as PbExprType;
163
164        let input_expr = InputRefExpression::new(order_data_type.clone(), 0);
165        let gap_expr = LiteralExpression::new(gap_data_type.clone(), Some(self.gap.clone()));
166        let add_expr = build_func(
167            PbExprType::Add,
168            order_data_type.clone(),
169            vec![input_expr.boxed(), gap_expr.boxed()],
170        )?;
171        let crate::expr::BoxedExpression::Sync(add_expr) = add_expr else {
172            bail!("session frame gap add expression must be sync");
173        };
174        self.add_expr = Some(add_expr);
175        Ok(())
176    }
177
178    pub fn new_for_test(
179        gap: ScalarImpl,
180        order_data_type: &DataType,
181        gap_data_type: &DataType,
182    ) -> Self {
183        let mut gap = Self::new(gap);
184        gap.prepare(order_data_type, gap_data_type).unwrap();
185        gap
186    }
187
188    fn for_calc(&self) -> SessionFrameGapRef<'_> {
189        SessionFrameGapRef {
190            add_expr: self.add_expr.as_ref().unwrap().as_ref(),
191        }
192    }
193}
194
195#[derive(Debug, Educe)]
196#[educe(Clone, Copy)]
197struct SessionFrameGapRef<'a> {
198    add_expr: &'a dyn SyncExpression,
199}
200
201impl SessionFrameGapRef<'_> {
202    fn minimal_next_start_of(&self, end_order_value: impl ToOwnedDatum) -> Datum {
203        let row = OwnedRow::new(vec![end_order_value.to_owned_datum()]);
204        self.add_expr
205            .eval_row(&row)
206            .expect("just simple calculation, should succeed") // TODO(rc): handle overflow
207    }
208}
risingwave_expr/window_function/session.rs

risingwave_expr/window_function/
session.rs
