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