Skip to main content

risingwave_frontend/expr/
correlated_input_ref.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 core::fmt;
16
17use risingwave_common::types::DataType;
18
19use super::{Expr, ExprImpl, InputRef};
20use crate::expr::ExprRewriter;
21
22pub type CorrelatedId = u32;
23
24/// Identifies a binder scope: a [`BindContext`](crate::binder::BindContext) whose columns can be
25/// referenced from a nested scope.
26#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
27pub struct ScopeId(pub(crate) u32);
28
29/// `Scope` is the binder scope that owns the referenced column.
30/// `Absolute` is the id of the Apply operator providing that scope's columns, and is non-zero.
31#[derive(Clone, Eq, PartialEq, Hash)]
32pub enum Position {
33    Scope(ScopeId),
34    Absolute(CorrelatedId),
35}
36
37/// A reference to a column outside the subquery.
38///
39/// `index` is the index in the referred relation.
40/// `position` has two modes: the binder records the owning scope, and the planner replaces it
41/// with the id of the Apply operator that binds that scope.
42#[derive(Clone, Eq, PartialEq, Hash)]
43pub struct CorrelatedInputRef {
44    index: usize,
45    data_type: DataType,
46    position: Position,
47}
48
49impl CorrelatedInputRef {
50    pub fn new(index: usize, data_type: DataType, scope: ScopeId) -> Self {
51        CorrelatedInputRef {
52            index,
53            data_type,
54            position: Position::Scope(scope),
55        }
56    }
57
58    /// Get a reference to the input ref's index.
59    pub fn index(&self) -> usize {
60        self.index
61    }
62
63    /// The scope owning the referenced column, if no Apply operator has been assigned yet.
64    pub fn scope(&self) -> Option<ScopeId> {
65        match self.position {
66            Position::Scope(scope) => Some(scope),
67            Position::Absolute(_) => None,
68        }
69    }
70
71    pub fn set_correlated_id(&mut self, correlated_id: CorrelatedId) {
72        self.position = Position::Absolute(correlated_id);
73    }
74
75    pub fn correlated_id(&self) -> CorrelatedId {
76        match self.position {
77            Position::Scope(_) => 0,
78            Position::Absolute(correlated_id) => correlated_id,
79        }
80    }
81}
82
83impl Expr for CorrelatedInputRef {
84    fn return_type(&self) -> DataType {
85        self.data_type.clone()
86    }
87
88    fn try_to_expr_proto(&self) -> Result<risingwave_pb::expr::ExprNode, String> {
89        Err(format!(
90            "CorrelatedInputRef {:?} has not been decorrelated",
91            self
92        ))
93    }
94}
95
96impl fmt::Debug for CorrelatedInputRef {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        f.debug_struct("CorrelatedInputRef")
99            .field("index", &self.index)
100            .field("correlated_id", &self.correlated_id())
101            .finish()
102    }
103}
104
105/// Rewrite `InputRef`s to columns of `scope` into `CorrelatedInputRef`s owned by it.
106pub struct InputRefScopeRewriter {
107    scope: ScopeId,
108}
109
110impl InputRefScopeRewriter {
111    pub fn new(scope: ScopeId) -> Self {
112        Self { scope }
113    }
114}
115
116impl ExprRewriter for InputRefScopeRewriter {
117    fn rewrite_input_ref(&mut self, input_ref: InputRef) -> ExprImpl {
118        CorrelatedInputRef::new(input_ref.index(), input_ref.return_type(), self.scope).into()
119    }
120}