risingwave_frontend/optimizer/plan_node/generic/
hop_window.rs

1// Copyright 2022 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::num::NonZeroUsize;
16
17use itertools::Itertools;
18use pretty_xmlish::{Pretty, StrAssocArr};
19use risingwave_common::catalog::{Field, Schema};
20use risingwave_common::types::{DataType, Interval};
21use risingwave_common::util::column_index_mapping::ColIndexMapping;
22use risingwave_expr::ExprError;
23
24use super::super::utils::IndicesDisplay;
25use super::{GenericPlanNode, GenericPlanRef, impl_distill_unit_from_fields};
26use crate::error::Result;
27use crate::expr::{ExprImpl, ExprType, FunctionCall, InputRef, InputRefDisplay, Literal};
28use crate::optimizer::optimizer_context::OptimizerContextRef;
29use crate::optimizer::property::FunctionalDependencySet;
30use crate::utils::ColIndexMappingRewriteExt;
31
32/// [`HopWindow`] implements Hop Table Function.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct HopWindow<PlanRef> {
35    pub input: PlanRef,
36    pub time_col: InputRef,
37    pub window_slide: Interval,
38    pub window_size: Interval,
39    pub window_offset: Interval,
40    /// Provides mapping from input schema, `window_start`, `window_end` to output schema.
41    /// For example, if we had:
42    /// input schema: | 0: `trip_time` | 1: `trip_name` |
43    /// `window_start`: 2
44    /// `window_end`: 3
45    /// output schema: | `trip_name` | `window_start` |
46    /// Then, `output_indices`: [1, 2]
47    pub output_indices: Vec<usize>,
48}
49
50impl<PlanRef: GenericPlanRef> GenericPlanNode for HopWindow<PlanRef> {
51    fn schema(&self) -> Schema {
52        let output_type = DataType::window_of(&self.time_col.data_type).unwrap();
53        let mut original_schema = self.input.schema().clone();
54        original_schema.fields.reserve_exact(2);
55        let window_start = Field::with_name(output_type.clone(), "window_start");
56        let window_end = Field::with_name(output_type, "window_end");
57        original_schema.fields.push(window_start);
58        original_schema.fields.push(window_end);
59        self.output_indices
60            .iter()
61            .map(|&idx| original_schema[idx].clone())
62            .collect()
63    }
64
65    fn stream_key(&self) -> Option<Vec<usize>> {
66        let window_start_index = self
67            .output_indices
68            .iter()
69            .position(|&idx| idx == self.input.schema().len());
70        let window_end_index = self
71            .output_indices
72            .iter()
73            .position(|&idx| idx == self.input.schema().len() + 1);
74        if window_start_index.is_none() && window_end_index.is_none() {
75            None
76        } else {
77            let mut pk = self
78                .input
79                .stream_key()?
80                .iter()
81                .filter_map(|&pk_idx| self.output_indices.iter().position(|&idx| idx == pk_idx))
82                .collect_vec();
83            if let Some(start_idx) = window_start_index {
84                pk.push(start_idx);
85            };
86            if let Some(end_idx) = window_end_index {
87                pk.push(end_idx);
88            };
89            Some(pk)
90        }
91    }
92
93    fn ctx(&self) -> OptimizerContextRef {
94        self.input.ctx()
95    }
96
97    fn functional_dependency(&self) -> FunctionalDependencySet {
98        let mut fd_set = self
99            .i2o_col_mapping()
100            .rewrite_functional_dependency_set(self.input.functional_dependency().clone());
101        let (start_idx_in_output, end_idx_in_output) = {
102            let internal2output = self.internal2output_col_mapping();
103            (
104                internal2output.try_map(self.internal_window_start_col_idx()),
105                internal2output.try_map(self.internal_window_end_col_idx()),
106            )
107        };
108        if let Some(start_idx) = start_idx_in_output
109            && let Some(end_idx) = end_idx_in_output
110        {
111            fd_set.add_functional_dependency_by_column_indices(&[start_idx], &[end_idx]);
112            fd_set.add_functional_dependency_by_column_indices(&[end_idx], &[start_idx]);
113        }
114        fd_set
115    }
116}
117
118impl<PlanRef: GenericPlanRef> HopWindow<PlanRef> {
119    pub fn clone_with_input<OtherPlanRef>(&self, input: OtherPlanRef) -> HopWindow<OtherPlanRef> {
120        HopWindow {
121            input,
122            time_col: self.time_col.clone(),
123            window_slide: self.window_slide,
124            window_size: self.window_size,
125            window_offset: self.window_offset,
126            output_indices: self.output_indices.clone(),
127        }
128    }
129
130    pub fn output_window_start_col_idx(&self) -> Option<usize> {
131        self.internal2output_col_mapping()
132            .try_map(self.internal_window_start_col_idx())
133    }
134
135    pub fn output_window_end_col_idx(&self) -> Option<usize> {
136        self.internal2output_col_mapping()
137            .try_map(self.internal_window_end_col_idx())
138    }
139
140    pub fn into_parts(self) -> (PlanRef, InputRef, Interval, Interval, Interval, Vec<usize>) {
141        (
142            self.input,
143            self.time_col,
144            self.window_slide,
145            self.window_size,
146            self.window_offset,
147            self.output_indices,
148        )
149    }
150
151    pub fn internal_window_start_col_idx(&self) -> usize {
152        self.input.schema().len()
153    }
154
155    pub fn internal_window_end_col_idx(&self) -> usize {
156        self.input.schema().len() + 1
157    }
158
159    pub fn o2i_col_mapping(&self) -> ColIndexMapping {
160        self.output2internal_col_mapping()
161            .composite(&self.internal2input_col_mapping())
162    }
163
164    pub fn i2o_col_mapping(&self) -> ColIndexMapping {
165        self.input2internal_col_mapping()
166            .composite(&self.internal2output_col_mapping())
167    }
168
169    pub fn internal_column_num(&self) -> usize {
170        self.input.schema().len() + 2
171    }
172
173    pub fn output2internal_col_mapping(&self) -> ColIndexMapping {
174        self.internal2output_col_mapping()
175            .inverse()
176            .expect("must be invertible")
177    }
178
179    pub fn internal2output_col_mapping(&self) -> ColIndexMapping {
180        ColIndexMapping::with_remaining_columns(&self.output_indices, self.internal_column_num())
181    }
182
183    pub fn input2internal_col_mapping(&self) -> ColIndexMapping {
184        ColIndexMapping::identity_or_none(self.input.schema().len(), self.internal_column_num())
185    }
186
187    pub fn internal2input_col_mapping(&self) -> ColIndexMapping {
188        ColIndexMapping::identity_or_none(self.internal_column_num(), self.input.schema().len())
189    }
190
191    pub fn derive_window_start_and_end_exprs(&self) -> Result<(Vec<ExprImpl>, Vec<ExprImpl>)> {
192        let Self {
193            window_size,
194            window_slide,
195            window_offset,
196            time_col,
197            ..
198        } = &self;
199        let units = window_size
200            .exact_div(window_slide)
201            .and_then(|x| NonZeroUsize::new(usize::try_from(x).ok()?))
202            .ok_or_else(|| ExprError::InvalidParam {
203                name: "window",
204                reason: format!(
205                    "window_size {} cannot be divided by window_slide {}",
206                    window_size, window_slide
207                )
208                .into(),
209            })?
210            .get();
211        let window_size_expr: ExprImpl =
212            Literal::new(Some((*window_size).into()), DataType::Interval).into();
213        let window_slide_expr: ExprImpl =
214            Literal::new(Some((*window_slide).into()), DataType::Interval).into();
215        let window_offset_expr: ExprImpl =
216            Literal::new(Some((*window_offset).into()), DataType::Interval).into();
217
218        let window_size_sub_slide = FunctionCall::new(
219            ExprType::Subtract,
220            vec![window_size_expr, window_slide_expr.clone()],
221        )?
222        .into();
223
224        let time_col_shifted = FunctionCall::new(
225            ExprType::Subtract,
226            vec![
227                ExprImpl::InputRef(Box::new(time_col.clone())),
228                window_size_sub_slide,
229            ],
230        )?
231        .into();
232
233        let hop_start: ExprImpl = FunctionCall::new(
234            ExprType::TumbleStart,
235            vec![time_col_shifted, window_slide_expr, window_offset_expr],
236        )?
237        .into();
238
239        let mut window_start_exprs = Vec::with_capacity(units);
240        let mut window_end_exprs = Vec::with_capacity(units);
241        for i in 0..units {
242            {
243                let window_start_offset =
244                    window_slide
245                        .checked_mul_int(i)
246                        .ok_or_else(|| ExprError::InvalidParam {
247                            name: "window",
248                            reason: format!(
249                                "window_slide {} cannot be multiplied by {}",
250                                window_slide, i
251                            )
252                            .into(),
253                        })?;
254                let window_start_offset_expr =
255                    Literal::new(Some(window_start_offset.into()), DataType::Interval).into();
256                let window_start_expr = FunctionCall::new(
257                    ExprType::Add,
258                    vec![hop_start.clone(), window_start_offset_expr],
259                )?
260                .into();
261                window_start_exprs.push(window_start_expr);
262            }
263            {
264                let window_end_offset =
265                    window_slide.checked_mul_int(i + units).ok_or_else(|| {
266                        ExprError::InvalidParam {
267                            name: "window",
268                            reason: format!(
269                                "window_slide {} cannot be multiplied by {}",
270                                window_slide,
271                                i + units
272                            )
273                            .into(),
274                        }
275                    })?;
276                let window_end_offset_expr =
277                    Literal::new(Some(window_end_offset.into()), DataType::Interval).into();
278                let window_end_expr = FunctionCall::new(
279                    ExprType::Add,
280                    vec![hop_start.clone(), window_end_offset_expr],
281                )?
282                .into();
283                window_end_exprs.push(window_end_expr);
284            }
285        }
286        assert_eq!(window_start_exprs.len(), window_end_exprs.len());
287        Ok((window_start_exprs, window_end_exprs))
288    }
289
290    pub fn fields_pretty<'a>(&self) -> StrAssocArr<'a> {
291        let mut out = Vec::with_capacity(5);
292        let output_type = DataType::window_of(&self.time_col.data_type).unwrap();
293        out.push((
294            "time_col",
295            Pretty::display(&InputRefDisplay {
296                input_ref: &self.time_col,
297                input_schema: self.input.schema(),
298            }),
299        ));
300        out.push(("slide", Pretty::display(&self.window_slide)));
301        out.push(("size", Pretty::display(&self.window_size)));
302        if self
303            .output_indices
304            .iter()
305            .copied()
306            // Behavior is the same as `LogicalHopWindow::internal_column_num`
307            .eq(0..(self.input.schema().len() + 2))
308        {
309            out.push(("output", Pretty::from("all")));
310        } else {
311            let original_schema: Schema = self
312                .input
313                .schema()
314                .clone()
315                .into_fields()
316                .into_iter()
317                .chain([
318                    Field::with_name(output_type.clone(), "window_start"),
319                    Field::with_name(output_type, "window_end"),
320                ])
321                .collect();
322            let id = IndicesDisplay {
323                indices: &self.output_indices,
324                schema: &original_schema,
325            };
326            out.push(("output", id.distill()));
327        }
328        out
329    }
330}
331
332impl_distill_unit_from_fields!(HopWindow, GenericPlanRef);
risingwave_frontend/optimizer/plan_node/generic/hop_window.rs

risingwave_frontend/optimizer/plan_node/generic/
hop_window.rs
