Skip to main content

risingwave_frontend/binder/
query.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::cell::RefCell;
16use std::collections::{HashMap, HashSet};
17use std::rc::Rc;
18
19use risingwave_common::catalog::Schema;
20use risingwave_common::types::DataType;
21use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
22use risingwave_sqlparser::ast::{
23    Cte, CteInner, Expr, Fetch, Ident, OrderByExpr, Query, Value, With,
24};
25use thiserror_ext::AsReport;
26
27use super::BoundValues;
28use super::bind_context::BindingCteState;
29use super::statement::RewriteExprsRecursive;
30use crate::binder::bind_context::BindingCte;
31use crate::binder::{Binder, BoundSetExpr};
32use crate::error::{ErrorCode, Result, RwError};
33use crate::expr::{CorrelatedId, Depth, ExprImpl, ExprRewriter};
34
35/// A validated sql query, including order and union.
36/// An example of its relationship with `BoundSetExpr` and `BoundSelect` can be found here: <https://bit.ly/3GQwgPz>
37#[derive(Debug, Clone)]
38pub struct BoundQuery {
39    pub body: BoundSetExpr,
40    pub order: Vec<ColumnOrder>,
41    pub limit: Option<u64>,
42    pub offset: Option<u64>,
43    pub with_ties: bool,
44    pub extra_order_exprs: Vec<ExprImpl>,
45}
46
47impl BoundQuery {
48    /// The schema returned by this [`BoundQuery`].
49    pub fn schema(&self) -> std::borrow::Cow<'_, Schema> {
50        self.body.schema()
51    }
52
53    /// The types returned by this [`BoundQuery`].
54    pub fn data_types(&self) -> Vec<DataType> {
55        self.schema().data_types()
56    }
57
58    /// Checks whether this query contains references to outer queries.
59    ///
60    /// Note there are 3 cases:
61    /// ```sql
62    /// select 1 from a having exists ( -- this is self
63    ///   select 1 from b where exists (
64    ///     select b1 from c
65    ///   )
66    /// );
67    ///
68    /// select 1 from a having exists ( -- this is self
69    ///   select 1 from b where exists (
70    ///     select a1 from c
71    ///   )
72    /// );
73    ///
74    /// select 1 from a where exists (
75    ///   select 1 from b having exists ( -- this is self, not the one above
76    ///     select a1 from c
77    ///   )
78    /// );
79    /// ```
80    /// We assume `self` is the subquery after `having`. In other words, the query with `from b` in
81    /// first 2 examples and the query with `from c` in the last example.
82    ///
83    /// * The first example is uncorrelated, because it is self-contained and does not depend on
84    ///   table `a`, although there is correlated input ref (`b1`) in it.
85    /// * The second example is correlated, because it depend on a correlated input ref (`a1`) that
86    ///   goes out.
87    /// * The last example is also correlated. because it cannot be evaluated independently either.
88    pub fn is_correlated_by_depth(&self, depth: Depth) -> bool {
89        self.body.is_correlated_by_depth(depth + 1)
90            || self
91                .extra_order_exprs
92                .iter()
93                .any(|e| e.has_correlated_input_ref_by_depth(depth + 1))
94    }
95
96    pub fn is_correlated_by_correlated_id(&self, correlated_id: CorrelatedId) -> bool {
97        self.body.is_correlated_by_correlated_id(correlated_id)
98            || self
99                .extra_order_exprs
100                .iter()
101                .any(|e| e.has_correlated_input_ref_by_correlated_id(correlated_id))
102    }
103
104    pub fn collect_correlated_indices_by_depth_and_assign_id(
105        &mut self,
106        depth: Depth,
107        correlated_id: CorrelatedId,
108    ) -> Vec<usize> {
109        let mut correlated_indices = vec![];
110
111        correlated_indices.extend(
112            self.body
113                .collect_correlated_indices_by_depth_and_assign_id(depth + 1, correlated_id),
114        );
115
116        correlated_indices.extend(self.extra_order_exprs.iter_mut().flat_map(|expr| {
117            expr.collect_correlated_indices_by_depth_and_assign_id(depth + 1, correlated_id)
118        }));
119        correlated_indices
120    }
121
122    /// Simple `VALUES` without other clauses.
123    pub fn with_values(values: BoundValues) -> Self {
124        BoundQuery {
125            body: BoundSetExpr::Values(values.into()),
126            order: vec![],
127            limit: None,
128            offset: None,
129            with_ties: false,
130            extra_order_exprs: vec![],
131        }
132    }
133}
134
135impl RewriteExprsRecursive for BoundQuery {
136    fn rewrite_exprs_recursive(&mut self, rewriter: &mut impl ExprRewriter) {
137        let new_extra_order_exprs = std::mem::take(&mut self.extra_order_exprs)
138            .into_iter()
139            .map(|expr| rewriter.rewrite_expr(expr))
140            .collect::<Vec<_>>();
141        self.extra_order_exprs = new_extra_order_exprs;
142
143        self.body.rewrite_exprs_recursive(rewriter);
144    }
145}
146
147impl Binder {
148    /// Bind a [`Query`].
149    ///
150    /// Before binding the [`Query`], we push the current [`BindContext`](super::BindContext) to the
151    /// stack and create a new context, because it may be a subquery.
152    ///
153    /// After finishing binding, we pop the previous context from the stack.
154    pub fn bind_query(&mut self, query: &Query) -> Result<BoundQuery> {
155        self.push_context();
156        let result = self.bind_query_inner(query);
157        self.pop_context()?;
158        result
159    }
160
161    /// Bind a [`Query`] for view.
162    /// TODO: support `SECURITY INVOKER` for view.
163    pub fn bind_query_for_view(&mut self, query: &Query) -> Result<BoundQuery> {
164        self.push_context();
165        self.context.disable_security_invoker = true;
166        let result = self.bind_query_inner(query);
167        self.pop_context()?;
168        result
169    }
170
171    /// Bind a [`Query`] using the current [`BindContext`](super::BindContext).
172    pub(super) fn bind_query_inner(
173        &mut self,
174        Query {
175            with,
176            body,
177            order_by,
178            limit,
179            offset,
180            fetch,
181        }: &Query,
182    ) -> Result<BoundQuery> {
183        let mut with_ties = false;
184        let limit = match (limit, fetch) {
185            (None, None) => None,
186            (
187                None,
188                Some(Fetch {
189                    with_ties: fetch_with_ties,
190                    quantity,
191                }),
192            ) => {
193                with_ties = *fetch_with_ties;
194                match quantity {
195                    Some(v) => Some(Expr::Value(Value::Number(v.clone()))),
196                    None => Some(Expr::Value(Value::Number("1".to_owned()))),
197                }
198            }
199            (Some(limit), None) => Some(limit.clone()),
200            (Some(_), Some(_)) => unreachable!(), // parse error
201        };
202        let limit_expr = limit.map(|expr| self.bind_expr(&expr)).transpose()?;
203        let limit = if let Some(limit_expr) = limit_expr {
204            // wrong type error is handled here
205            let limit_cast_to_bigint = limit_expr.cast_assign(&DataType::Int64).map_err(|_| {
206                RwError::from(ErrorCode::ExprError(
207                    "expects an integer or expression that can be evaluated to an integer after LIMIT"
208                        .into(),
209                ))
210            })?;
211            let limit = match limit_cast_to_bigint.try_fold_const() {
212                Some(Ok(Some(datum))) => {
213                    let value = datum.as_int64();
214                    if *value < 0 {
215                        return Err(ErrorCode::ExprError(
216                            format!("LIMIT must not be negative, but found: {}", *value).into(),
217                        )
218                            .into());
219                    }
220                    *value as u64
221                }
222                // If evaluated to NULL, we follow PG to treat NULL as no limit
223                Some(Ok(None)) => {
224                    u64::MAX
225                }
226                // not const error
227                None => return Err(ErrorCode::ExprError(
228                    "expects an integer or expression that can be evaluated to an integer after LIMIT, but found non-const expression"
229                        .into(),
230                ).into()),
231                // eval error
232                Some(Err(e)) => {
233                    return Err(ErrorCode::ExprError(
234                        format!("expects an integer or expression that can be evaluated to an integer after LIMIT,\nbut the evaluation of the expression returns error:{}", e.as_report()
235                        ).into(),
236                    ).into())
237                }
238            };
239            Some(limit)
240        } else {
241            None
242        };
243
244        let offset = offset
245            .as_ref()
246            .map(|s| parse_non_negative_i64("OFFSET", s))
247            .transpose()?
248            .map(|v| v as u64);
249
250        if let Some(with) = with {
251            self.bind_with(with)?;
252        }
253        let body = self.bind_set_expr(body)?;
254        let name_to_index =
255            Self::build_name_to_index(body.schema().fields().iter().map(|f| f.name.clone()));
256        let mut extra_order_exprs = vec![];
257        let visible_output_num = body.schema().len();
258        let order = order_by
259            .iter()
260            .map(|order_by_expr| {
261                self.bind_order_by_expr_in_query(
262                    order_by_expr,
263                    &body,
264                    &name_to_index,
265                    &mut extra_order_exprs,
266                    visible_output_num,
267                )
268            })
269            .collect::<Result<_>>()?;
270        Ok(BoundQuery {
271            body,
272            order,
273            limit,
274            offset,
275            with_ties,
276            extra_order_exprs,
277        })
278    }
279
280    pub fn build_name_to_index(names: impl Iterator<Item = String>) -> HashMap<String, usize> {
281        let mut m = HashMap::new();
282        names.enumerate().for_each(|(index, name)| {
283            m.entry(name)
284                // Ambiguous (duplicate) output names are marked with usize::MAX.
285                // This is not necessarily an error as long as not actually referenced.
286                .and_modify(|v| *v = usize::MAX)
287                .or_insert(index);
288        });
289        m
290    }
291
292    /// Bind an `ORDER BY` expression in a [`Query`], which can be either:
293    /// * an output-column name
294    /// * index of an output column
295    /// * an arbitrary expression
296    ///
297    /// Refer to `bind_group_by_expr_in_select` to see their similarities and differences.
298    ///
299    /// # Arguments
300    ///
301    /// * `name_to_index` - visible output column name -> index. Ambiguous (duplicate) output names
302    ///   are marked with `usize::MAX`.
303    /// * `visible_output_num` - the number of all visible output columns, including duplicates.
304    fn bind_order_by_expr_in_query(
305        &mut self,
306        OrderByExpr {
307            expr,
308            asc,
309            nulls_first,
310        }: &OrderByExpr,
311        body: &BoundSetExpr,
312        name_to_index: &HashMap<String, usize>,
313        extra_order_exprs: &mut Vec<ExprImpl>,
314        visible_output_num: usize,
315    ) -> Result<ColumnOrder> {
316        let order_type = OrderType::from_bools(*asc, *nulls_first);
317
318        // If the query body is a simple `SELECT`, we can reuse an existing select item by
319        // expression equality, instead of always appending a new hidden column for ORDER BY.
320        //
321        // This is only safe for pure expressions. For example, `ORDER BY random()` must not be
322        // rewritten to reuse `SELECT random()` because they should be evaluated independently.
323        let select_items_for_match = match body {
324            BoundSetExpr::Select(s) => Some(&s.select_items[..]),
325            _ => None,
326        };
327
328        let column_index = match expr {
329            Expr::Identifier(name) if let Some(index) = name_to_index.get(&name.real_value()) => {
330                match *index != usize::MAX {
331                    true => *index,
332                    false => {
333                        return Err(ErrorCode::BindError(format!(
334                            "ORDER BY \"{}\" is ambiguous",
335                            name.real_value()
336                        ))
337                        .into());
338                    }
339                }
340            }
341            Expr::Value(Value::Number(number)) => match number.parse::<usize>() {
342                Ok(index) if 1 <= index && index <= visible_output_num => index - 1,
343                _ => {
344                    return Err(ErrorCode::InvalidInputSyntax(format!(
345                        "Invalid ordinal number in ORDER BY: {}",
346                        number
347                    ))
348                    .into());
349                }
350            },
351            expr => {
352                let bound_expr = self.bind_expr(expr)?;
353
354                // Subqueries are unnested later, and comparing expressions that still contain
355                // subqueries is intentionally unsupported.
356                if bound_expr.is_pure()
357                    && !bound_expr.has_subquery()
358                    && let Some(select_items) = select_items_for_match
359                    && let Some(existing_idx) = select_items.iter().position(|e| e == &bound_expr)
360                {
361                    existing_idx
362                } else {
363                    extra_order_exprs.push(bound_expr);
364                    visible_output_num + extra_order_exprs.len() - 1
365                }
366            }
367        };
368        Ok(ColumnOrder::new(column_index, order_type))
369    }
370
371    fn bind_with(&mut self, with: &With) -> Result<()> {
372        if with.recursive {
373            return Err(ErrorCode::BindError("RECURSIVE CTE is not supported".to_owned()).into());
374        }
375
376        let mut cte_names: HashSet<String> = HashSet::new();
377
378        for cte_table in &with.cte_tables {
379            let share_id = self.next_share_id();
380            let Cte { alias, cte_inner } = cte_table;
381            let table_name = alias.name.real_value();
382
383            // Check whether there are duplicate CTE names within the same WITH clause.
384            if cte_names.contains(&table_name) {
385                return Err(ErrorCode::DuplicateRelationName(format!(
386                    "WITH query name \"{}\" specified more than once",
387                    table_name
388                ))
389                .into());
390            }
391
392            cte_names.insert(table_name.clone());
393
394            match cte_inner {
395                CteInner::Query(query) => {
396                    let bound_query = self.bind_query(query)?;
397                    self.context.cte_to_relation.insert(
398                        table_name,
399                        Rc::new(RefCell::new(BindingCte {
400                            share_id,
401                            state: BindingCteState::Bound { query: bound_query },
402                            alias: alias.clone(),
403                        })),
404                    );
405                }
406                CteInner::ChangeLog { from, key } => {
407                    self.push_context();
408                    let from_table_relation = self.bind_relation_by_name(from, None, None, true)?;
409                    let key_indices = key
410                        .as_ref()
411                        .map(|key| self.bind_changelog_key(key))
412                        .transpose()?;
413
414                    self.pop_context()?;
415                    self.context.cte_to_relation.insert(
416                        table_name,
417                        Rc::new(RefCell::new(BindingCte {
418                            share_id,
419                            state: BindingCteState::ChangeLog {
420                                table: from_table_relation,
421                                key_indices,
422                            },
423                            alias: alias.clone(),
424                        })),
425                    );
426                }
427            }
428        }
429        Ok(())
430    }
431
432    fn bind_changelog_key(&self, key: &[Ident]) -> Result<Vec<usize>> {
433        let mut seen = HashSet::new();
434
435        key.iter()
436            .map(|ident| {
437                let name = ident.real_value();
438                let index = self.context.get_column_binding_index(&None, &None, &name)?;
439                let column = &self.context.columns[index];
440
441                if column.is_hidden {
442                    return Err(ErrorCode::BindError(format!(
443                        "Changelog key column \"{name}\" must not be hidden"
444                    ))
445                    .into());
446                }
447
448                if !seen.insert(column.index) {
449                    return Err(ErrorCode::BindError(format!(
450                        "Changelog key column \"{name}\" specified more than once"
451                    ))
452                    .into());
453                }
454
455                Ok(index)
456            })
457            .collect()
458    }
459}
460
461// TODO: Make clause a const generic param after <https://github.com/rust-lang/rust/issues/95174>.
462fn parse_non_negative_i64(clause: &str, s: &str) -> Result<i64> {
463    match s.parse::<i64>() {
464        Ok(v) => {
465            if v < 0 {
466                Err(ErrorCode::InvalidInputSyntax(format!("{clause} must not be negative")).into())
467            } else {
468                Ok(v)
469            }
470        }
471        Err(e) => Err(ErrorCode::InvalidInputSyntax(e.to_report_string()).into()),
472    }
473}