risingwave_frontend/expr/
subquery.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::hash::Hash;
16
17use risingwave_common::types::{DataType, StructType};
18
19use super::{Expr, ExprImpl, ExprType};
20use crate::binder::{BoundQuery, UNNAMED_COLUMN};
21use crate::expr::{CorrelatedId, ScopeId};
22
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub enum SubqueryKind {
25    /// Returns a scalar value (single column single row).
26    Scalar,
27    /// Returns a scalar struct value composed of multiple columns.
28    /// Used in `UPDATE SET (col1, col2) = (SELECT ...)`.
29    UpdateSet,
30    /// `EXISTS` | `NOT EXISTS` subquery (semi/anti-semi join). Returns a boolean.
31    Existential,
32    /// `IN` subquery.
33    In(ExprImpl),
34    /// Expression operator `SOME` subquery.
35    Some(ExprImpl, ExprType),
36    /// Expression operator `ALL` subquery.
37    All(ExprImpl, ExprType),
38    /// Expression operator `ARRAY` subquery.
39    Array,
40}
41
42/// Subquery expression.
43#[derive(Clone)]
44pub struct Subquery {
45    pub query: BoundQuery,
46    pub kind: SubqueryKind,
47    /// The scope this subquery is evaluated in. Its references to columns of this scope are bound
48    /// to the Apply operator that plans the subquery.
49    pub outer_scope: ScopeId,
50}
51
52impl Subquery {
53    pub fn new(query: BoundQuery, kind: SubqueryKind, outer_scope: ScopeId) -> Self {
54        Self {
55            query,
56            kind,
57            outer_scope,
58        }
59    }
60
61    pub fn is_correlated_by_scope(&self, scope: ScopeId) -> bool {
62        let is_correlated = match &self.kind {
63            SubqueryKind::In(expr) => expr.has_correlated_input_ref_by_scope(scope),
64            SubqueryKind::Some(expr, _) | SubqueryKind::All(expr, _) => {
65                expr.has_correlated_input_ref_by_scope(scope)
66            }
67            SubqueryKind::Array
68            | SubqueryKind::Scalar
69            | SubqueryKind::UpdateSet
70            | SubqueryKind::Existential => false,
71        };
72        is_correlated || self.query.is_correlated_by_scope(scope)
73    }
74
75    pub fn is_correlated_by_correlated_id(&self, correlated_id: CorrelatedId) -> bool {
76        let is_correlated = match &self.kind {
77            SubqueryKind::In(expr) => expr.has_correlated_input_ref_by_correlated_id(correlated_id),
78            SubqueryKind::Some(expr, _) | SubqueryKind::All(expr, _) => {
79                expr.has_correlated_input_ref_by_correlated_id(correlated_id)
80            }
81            SubqueryKind::Array
82            | SubqueryKind::Scalar
83            | SubqueryKind::UpdateSet
84            | SubqueryKind::Existential => false,
85        };
86        is_correlated || self.query.is_correlated_by_correlated_id(correlated_id)
87    }
88
89    pub fn collect_correlated_indices_by_scope_and_assign_id(
90        &mut self,
91        scope: ScopeId,
92        correlated_id: CorrelatedId,
93    ) -> Vec<usize> {
94        let mut correlated_indices = self
95            .query
96            .collect_correlated_indices_by_scope_and_assign_id(scope, correlated_id);
97
98        match &mut self.kind {
99            SubqueryKind::In(expr) => {
100                correlated_indices.extend(
101                    expr.collect_correlated_indices_by_scope_and_assign_id(scope, correlated_id),
102                );
103            }
104            SubqueryKind::Some(expr, _) | SubqueryKind::All(expr, _) => {
105                correlated_indices.extend(
106                    expr.collect_correlated_indices_by_scope_and_assign_id(scope, correlated_id),
107                );
108            }
109            SubqueryKind::Array
110            | SubqueryKind::Scalar
111            | SubqueryKind::UpdateSet
112            | SubqueryKind::Existential => {
113                // No additional correlated indices to collect for these kinds.
114            }
115        }
116        correlated_indices.sort();
117        correlated_indices.dedup();
118        correlated_indices
119    }
120}
121
122impl PartialEq for Subquery {
123    fn eq(&self, _other: &Self) -> bool {
124        unreachable!("Subquery {:?} has not been unnested", self)
125    }
126}
127
128impl Hash for Subquery {
129    fn hash<H: std::hash::Hasher>(&self, _state: &mut H) {
130        unreachable!("Subquery {:?} has not been hashed", self)
131    }
132}
133
134impl Eq for Subquery {}
135
136impl Expr for Subquery {
137    fn return_type(&self) -> DataType {
138        match self.kind {
139            SubqueryKind::Scalar => {
140                let types = self.query.data_types();
141                assert_eq!(types.len(), 1, "Subquery with more than one column");
142                types[0].clone()
143            }
144            SubqueryKind::UpdateSet => {
145                let schema = self.query.schema();
146                let struct_type = if schema.fields().iter().any(|f| f.name == UNNAMED_COLUMN) {
147                    StructType::unnamed(self.query.data_types())
148                } else {
149                    StructType::new(
150                        (schema.fields().iter().cloned()).map(|f| (f.name, f.data_type)),
151                    )
152                };
153                DataType::Struct(struct_type)
154            }
155            SubqueryKind::Array => {
156                let types = self.query.data_types();
157                assert_eq!(types.len(), 1, "Subquery with more than one column");
158                DataType::list(types[0].clone())
159            }
160            _ => DataType::Boolean,
161        }
162    }
163
164    fn try_to_expr_proto(&self) -> Result<risingwave_pb::expr::ExprNode, String> {
165        Err(format!("Subquery {:?} has not been unnested", self))
166    }
167}
168
169impl std::fmt::Debug for Subquery {
170    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171        f.debug_struct("Subquery")
172            .field("kind", &self.kind)
173            .field("query", &self.query)
174            .finish()
175    }
176}
risingwave_frontend/expr/subquery.rs

risingwave_frontend/expr/
subquery.rs
