risingwave_frontend/planner/
select.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::collections::HashMap;
16
17use itertools::Itertools;
18use risingwave_common::bail_not_implemented;
19use risingwave_common::catalog::PROJECTED_ROW_ID_COLUMN_NAME;
20use risingwave_common::types::DataType;
21use risingwave_common::util::iter_util::ZipEqFast;
22use risingwave_common::util::sort_util::ColumnOrder;
23use risingwave_expr::ExprError;
24use risingwave_pb::plan_common::JoinType;
25
26use crate::OptimizerContextRef;
27use crate::binder::{BoundDistinct, BoundSelect};
28use crate::error::{ErrorCode, Result};
29use crate::expr::{
30    CorrelatedId, Expr, ExprImpl, ExprRewriter, ExprType, FunctionCall, InputRef, Subquery,
31    SubqueryKind,
32};
33pub use crate::optimizer::plan_node::LogicalFilter;
34use crate::optimizer::plan_node::generic::{Agg, GenericPlanRef, Project, ProjectBuilder};
35use crate::optimizer::plan_node::{
36    LogicalAgg, LogicalApply, LogicalDedup, LogicalJoin, LogicalOverWindow,
37    LogicalPlanRef as PlanRef, LogicalProject, LogicalProjectSet, LogicalTopN, LogicalValues,
38    PlanAggCall,
39};
40use crate::optimizer::property::Order;
41use crate::planner::Planner;
42use crate::utils::{Condition, IndexSet};
43
44impl Planner {
45    pub(super) fn plan_select(
46        &mut self,
47        BoundSelect {
48            from,
49            where_clause,
50            mut select_items,
51            group_by,
52            mut having,
53            mut distinct,
54            ..
55        }: BoundSelect,
56        extra_order_exprs: Vec<ExprImpl>,
57        order: &[ColumnOrder],
58    ) -> Result<PlanRef> {
59        // Append expressions in ORDER BY.
60        if distinct.is_distinct() && !extra_order_exprs.is_empty() {
61            return Err(ErrorCode::InvalidInputSyntax(
62                "for SELECT DISTINCT, ORDER BY expressions must appear in select list".into(),
63            )
64            .into());
65        }
66        select_items.extend(extra_order_exprs);
67        // The DISTINCT ON expression(s) must match the leftmost ORDER BY expression(s).
68        if let BoundDistinct::DistinctOn(exprs) = &distinct {
69            let mut distinct_on_exprs: HashMap<ExprImpl, bool> =
70                exprs.iter().map(|expr| (expr.clone(), false)).collect();
71            let mut uncovered_distinct_on_exprs_cnt = distinct_on_exprs.len();
72            let mut order_iter = order.iter().map(|o| &select_items[o.column_index]);
73            while uncovered_distinct_on_exprs_cnt > 0
74                && let Some(order_expr) = order_iter.next()
75            {
76                match distinct_on_exprs.get_mut(order_expr) {
77                    Some(has_been_covered) => {
78                        if !*has_been_covered {
79                            *has_been_covered = true;
80                            uncovered_distinct_on_exprs_cnt -= 1;
81                        }
82                    }
83                    None => {
84                        return Err(ErrorCode::InvalidInputSyntax(
85                            "the SELECT DISTINCT ON expressions must match the leftmost ORDER BY expressions"
86                                .into(),
87                        )
88                        .into());
89                    }
90                }
91            }
92        }
93
94        // Plan the FROM clause.
95        let mut root = match from {
96            None => self.create_dummy_values(),
97            Some(t) => self.plan_relation(t)?,
98        };
99        // Plan the WHERE clause.
100        if let Some(where_clause) = where_clause {
101            root = self.plan_where(root, where_clause)?;
102        }
103        // Plan the SELECT clause.
104        // TODO: select-agg, group-by, having can also contain subquery exprs.
105        let has_agg_call = select_items.iter().any(|expr| expr.has_agg_call());
106        if !group_by.is_empty() || having.is_some() || has_agg_call {
107            // The `DISTINCT ON` expressions must also be rewritten to reference the
108            // aggregated results, just like the SELECT and ORDER BY expressions. By
109            // appending them to `select_items` they go through the same agg rewriting,
110            // which maps group-key references and rejects any column that neither
111            // appears in GROUP BY nor inside an aggregate (matching PostgreSQL).
112            let n_distinct_on = if let BoundDistinct::DistinctOn(exprs) = &distinct {
113                select_items.extend(exprs.iter().cloned());
114                exprs.len()
115            } else {
116                0
117            };
118
119            (root, select_items, having) =
120                LogicalAgg::create(select_items, group_by, having, root)?;
121
122            if n_distinct_on > 0 {
123                let rewritten = select_items.split_off(select_items.len() - n_distinct_on);
124                distinct = BoundDistinct::DistinctOn(rewritten);
125            }
126        }
127
128        if let Some(having) = having {
129            root = self.plan_where(root, having)?;
130        }
131
132        if select_items.iter().any(|e| e.has_subquery()) {
133            (root, select_items) =
134                self.substitute_subqueries_in_cross_join_way(root, select_items)?;
135        }
136        if select_items.iter().any(|e| e.has_window_function()) {
137            (root, select_items) = LogicalOverWindow::create(root, select_items)?;
138        }
139
140        let original_select_items_len = select_items.len();
141
142        // variable `distinct_list_index_to_select_items_index` is meaningful iff
143        // `matches!(&distinct, BoundDistinct::DistinctOn(_))`
144        let mut distinct_list_index_to_select_items_index = vec![];
145        if let BoundDistinct::DistinctOn(distinct_list) = &distinct {
146            distinct_list_index_to_select_items_index.reserve(distinct_list.len());
147            let mut builder_index_to_select_items_index =
148                Vec::with_capacity(original_select_items_len);
149            let mut input_proj_builder = ProjectBuilder::default();
150            for (select_item_index, select_item) in select_items.iter().enumerate() {
151                let builder_index = input_proj_builder
152                    .add_expr(select_item)
153                    .map_err(|msg| ExprError::UnsupportedFunction(String::from(msg)))?;
154                if builder_index >= builder_index_to_select_items_index.len() {
155                    debug_assert_eq!(builder_index, builder_index_to_select_items_index.len());
156                    builder_index_to_select_items_index.push(select_item_index);
157                }
158            }
159            for distinct_expr in distinct_list {
160                let builder_index = input_proj_builder
161                    .add_expr(distinct_expr)
162                    .map_err(|msg| ExprError::UnsupportedFunction(String::from(msg)))?;
163                if builder_index >= builder_index_to_select_items_index.len() {
164                    debug_assert_eq!(builder_index, builder_index_to_select_items_index.len());
165                    select_items.push(distinct_expr.clone());
166                    builder_index_to_select_items_index.push(select_items.len() - 1);
167                }
168                distinct_list_index_to_select_items_index
169                    .push(builder_index_to_select_items_index[builder_index]);
170            }
171        }
172
173        let need_restore_select_items = select_items.len() > original_select_items_len;
174
175        root = LogicalProjectSet::create(root, select_items);
176
177        if matches!(&distinct, BoundDistinct::DistinctOn(_)) {
178            root = if order.is_empty() {
179                // We only support deduplicating `DISTINCT ON` columns when there is no `ORDER BY`
180                // clause now.
181                LogicalDedup::new(root, distinct_list_index_to_select_items_index).into()
182            } else {
183                LogicalTopN::new(
184                    root,
185                    1,
186                    0,
187                    false,
188                    Order::new(order.to_vec()),
189                    distinct_list_index_to_select_items_index,
190                )
191                .into()
192            };
193        }
194
195        if need_restore_select_items {
196            root = LogicalProject::with_core(Project::with_out_col_idx(
197                root,
198                0..original_select_items_len,
199            ))
200            .into();
201        }
202
203        if let BoundDistinct::Distinct = distinct {
204            let fields = root.schema().fields();
205            let group_key = if let Some(field) = fields.first()
206                && field.name == PROJECTED_ROW_ID_COLUMN_NAME
207            {
208                // Do not group by projected_row_id hidden column.
209                (1..fields.len()).collect()
210            } else {
211                (0..fields.len()).collect()
212            };
213            root = Agg::new(vec![], group_key, root).into();
214        }
215
216        Ok(root)
217    }
218
219    /// Helper to create a dummy node as child of [`LogicalProject`].
220    /// For example, `select 1+2, 3*4` will be `Project([1+2, 3+4]) - Values([[]])`.
221    fn create_dummy_values(&self) -> PlanRef {
222        LogicalValues::create_empty_scalar(self.ctx.clone())
223    }
224
225    /// Helper to create an `EXISTS` boolean operator with the given `input`.
226    /// It is represented by `Project([$0 >= 1]) -> Agg(count(*)) -> input`
227    fn create_exists(&self, input: PlanRef) -> Result<PlanRef> {
228        let count_star = Agg::new(vec![PlanAggCall::count_star()], IndexSet::empty(), input);
229        let ge = FunctionCall::new(
230            ExprType::GreaterThanOrEqual,
231            vec![
232                InputRef::new(0, DataType::Int64).into(),
233                ExprImpl::literal_int(1),
234            ],
235        )
236        .unwrap();
237        Ok(LogicalProject::create(count_star.into(), vec![ge.into()]))
238    }
239
240    /// For `(NOT) EXISTS subquery` or `(NOT) IN subquery`, we can plan it as
241    /// `LeftSemi/LeftAnti` [`LogicalApply`]
242    /// For other subqueries, we plan it as `LeftOuter` [`LogicalApply`] using
243    /// [`Self::substitute_subqueries_in_left_deep_tree_way`].
244    pub(super) fn plan_where(
245        &mut self,
246        mut input: PlanRef,
247        where_clause: ExprImpl,
248    ) -> Result<PlanRef> {
249        if !where_clause.has_subquery() {
250            return Ok(LogicalFilter::create_with_expr(input, where_clause));
251        }
252        let (subquery_conjunctions, not_subquery_conjunctions, others) =
253            Condition::with_expr(where_clause)
254                .group_by::<_, 3>(|expr| match expr {
255                    ExprImpl::Subquery(_) => 0,
256                    ExprImpl::FunctionCall(func_call)
257                        if func_call.func_type() == ExprType::Not
258                            && matches!(func_call.inputs()[0], ExprImpl::Subquery(_)) =>
259                    {
260                        1
261                    }
262                    _ => 2,
263                })
264                .into_iter()
265                .next_tuple()
266                .unwrap();
267
268        // EXISTS and IN in WHERE.
269        for expr in subquery_conjunctions {
270            self.handle_exists_and_in(expr, false, &mut input)?;
271        }
272
273        // NOT EXISTS and NOT IN in WHERE.
274        for expr in not_subquery_conjunctions {
275            let not = expr.into_function_call().unwrap();
276            let (_, expr) = not.decompose_as_unary();
277            self.handle_exists_and_in(expr, true, &mut input)?;
278        }
279
280        if others.always_true() {
281            Ok(input)
282        } else {
283            let (input, others) =
284                self.substitute_subqueries_in_left_deep_tree_way(input, others.conjunctions)?;
285            Ok(LogicalFilter::create(
286                input,
287                Condition {
288                    conjunctions: others,
289                },
290            ))
291        }
292    }
293
294    /// Handle (NOT) EXISTS and (NOT) IN in WHERE clause.
295    ///
296    /// We will use a = b to replace a in (select b from ....) for (NOT) IN thus avoiding adding a
297    /// `LogicalFilter` on `LogicalApply`.
298    fn handle_exists_and_in(
299        &mut self,
300        expr: ExprImpl,
301        negated: bool,
302        input: &mut PlanRef,
303    ) -> Result<()> {
304        let join_type = if negated {
305            JoinType::LeftAnti
306        } else {
307            JoinType::LeftSemi
308        };
309        let correlated_id = self.ctx.next_correlated_id();
310        let mut subquery = expr.into_subquery().unwrap();
311        // we should call `subquery.query.collect_correlated_indices_by_scope_and_assign_id`
312        // instead of `subquery.collect_correlated_indices_by_scope_and_assign_id`.
313        // because current subquery containing struct `kind` expr which should never be correlated with the current subquery.
314        let mut correlated_indices = subquery
315            .query
316            .collect_correlated_indices_by_scope_and_assign_id(subquery.outer_scope, correlated_id);
317        correlated_indices.sort();
318        correlated_indices.dedup();
319        let output_column_type = subquery.query.data_types()[0].clone();
320        let right_plan = self.plan_query(subquery.query)?.into_unordered_subplan();
321        let on = match subquery.kind {
322            SubqueryKind::Existential => ExprImpl::literal_bool(true),
323            SubqueryKind::In(left_expr) => {
324                let right_expr = InputRef::new(input.schema().len(), output_column_type);
325                FunctionCall::new(ExprType::Equal, vec![left_expr, right_expr.into()])?.into()
326            }
327            kind => bail_not_implemented!(issue = 1343, "Not supported subquery kind: {:?}", kind),
328        };
329        *input = Self::create_apply(
330            correlated_id,
331            correlated_indices,
332            input.clone(),
333            right_plan,
334            on,
335            join_type,
336            false,
337        );
338        Ok(())
339    }
340
341    /// Substitutes all [`Subquery`] in `exprs` in a left deep tree way.
342    ///
343    /// Each time a [`Subquery`] is found, it is replaced by a new [`InputRef`]. And `root` is
344    /// replaced by a new `LeftOuter` [`LogicalApply`] whose left side is `root` and right side is
345    /// the planned subquery.
346    ///
347    /// The [`InputRef`]s' indexes start from `root.schema().len()`,
348    /// which means they are additional columns beyond the original `root`.
349    ///
350    /// The left-deep tree way only meaningful when there are multiple subqueries
351    ///
352    /// ```text
353    ///             Apply
354    ///            /    \
355    ///          Apply  Subquery3
356    ///         /    \
357    ///     Apply    Subquery2
358    ///     /   \
359    ///   left  Subquery1
360    /// ```
361    ///
362    /// Typically, this is used for subqueries in `WHERE` clause, because it is unlikely that users would write subqueries in the same where clause multiple times.
363    /// But our dynamic filter sometimes generates subqueries in the same where clause multiple times (which couldn't work in the cross join way), so we need to support this case.
364    pub(super) fn substitute_subqueries_in_left_deep_tree_way(
365        &mut self,
366        mut root: PlanRef,
367        mut exprs: Vec<ExprImpl>,
368    ) -> Result<(PlanRef, Vec<ExprImpl>)> {
369        struct SubstituteSubQueries {
370            input_col_num: usize,
371            subqueries: Vec<Subquery>,
372            correlated_indices_collection: Vec<Vec<usize>>,
373            correlated_ids: Vec<CorrelatedId>,
374            ctx: OptimizerContextRef,
375        }
376
377        // TODO: consider the multi-subquery case for normal predicate.
378        impl ExprRewriter for SubstituteSubQueries {
379            fn rewrite_subquery(&mut self, mut subquery: Subquery) -> ExprImpl {
380                let correlated_id = self.ctx.next_correlated_id();
381                self.correlated_ids.push(correlated_id);
382                let input_ref = InputRef::new(self.input_col_num, subquery.return_type()).into();
383                self.input_col_num += 1;
384                self.correlated_indices_collection.push(
385                    subquery.collect_correlated_indices_by_scope_and_assign_id(
386                        subquery.outer_scope,
387                        correlated_id,
388                    ),
389                );
390                self.subqueries.push(subquery);
391                input_ref
392            }
393        }
394
395        let mut rewriter = SubstituteSubQueries {
396            input_col_num: root.schema().len(),
397            subqueries: vec![],
398            correlated_indices_collection: vec![],
399            correlated_ids: vec![],
400            ctx: self.ctx.clone(),
401        };
402        exprs = exprs
403            .into_iter()
404            .map(|e| rewriter.rewrite_expr(e))
405            .collect();
406
407        for ((subquery, correlated_indices), correlated_id) in rewriter
408            .subqueries
409            .into_iter()
410            .zip_eq_fast(rewriter.correlated_indices_collection)
411            .zip_eq_fast(rewriter.correlated_ids)
412        {
413            let return_type = subquery.return_type();
414            let subroot = self.plan_query(subquery.query)?;
415
416            let right = match subquery.kind {
417                SubqueryKind::Scalar => subroot.into_unordered_subplan(),
418                SubqueryKind::UpdateSet => {
419                    let plan = subroot.into_unordered_subplan();
420
421                    // Compose all input columns into a struct with `ROW` function.
422                    let all_input_refs = plan
423                        .schema()
424                        .data_types()
425                        .into_iter()
426                        .enumerate()
427                        .map(|(i, data_type)| InputRef::new(i, data_type).into())
428                        .collect::<Vec<_>>();
429                    let call =
430                        FunctionCall::new_unchecked(ExprType::Row, all_input_refs, return_type);
431
432                    LogicalProject::create(plan, vec![call.into()])
433                }
434                SubqueryKind::Existential => {
435                    self.create_exists(subroot.into_unordered_subplan())?
436                }
437                SubqueryKind::Array => subroot.into_array_agg()?,
438                _ => bail_not_implemented!(issue = 1343, "{:?}", subquery.kind),
439            };
440
441            root = Self::create_apply(
442                correlated_id,
443                correlated_indices,
444                root,
445                right,
446                ExprImpl::literal_bool(true),
447                JoinType::LeftOuter,
448                true,
449            );
450        }
451        Ok((root, exprs))
452    }
453
454    /// Substitutes all [`Subquery`] in `exprs` in a cross join way.
455    ///
456    /// Each time a [`Subquery`] is found, it is replaced by a new [`InputRef`]. And `root` is
457    /// replaced by a new `LeftOuter` [`LogicalApply`] whose left side is `root` and right side is
458    /// the planned subquery.
459    ///
460    /// The [`InputRef`]s' indexes start from `root.schema().len()`,
461    /// which means they are additional columns beyond the original `root`.
462    ///
463    ///
464    /// The cross join way only meaningful when there are multiple subqueries
465    ///
466    /// ```text
467    ///            Apply
468    ///           /    \
469    ///         left   CrossJoin
470    ///                   /   \
471    ///           Subquery1   CrossJoin
472    ///                         /   \
473    ///                 Subquery2   Subquery3
474    /// ```
475    /// Typically, this is used for scalar subqueries in select clauses, because users might write subqueries in exprs multiple times.
476    /// If we use the left-deep tree way, it will generate a lot of `Apply` nodes, which is not efficient.
477    pub(super) fn substitute_subqueries_in_cross_join_way(
478        &mut self,
479        mut root: PlanRef,
480        mut exprs: Vec<ExprImpl>,
481    ) -> Result<(PlanRef, Vec<ExprImpl>)> {
482        struct SubstituteSubQueries {
483            input_col_num: usize,
484            subqueries: Vec<Subquery>,
485            correlated_id: Option<CorrelatedId>,
486            correlated_indices_collection: Vec<Vec<usize>>,
487            ctx: OptimizerContextRef,
488        }
489
490        impl ExprRewriter for SubstituteSubQueries {
491            fn rewrite_subquery(&mut self, mut subquery: Subquery) -> ExprImpl {
492                if self.correlated_id.is_none() {
493                    self.correlated_id = Some(self.ctx.next_correlated_id());
494                }
495                let input_ref = InputRef::new(self.input_col_num, subquery.return_type()).into();
496                self.input_col_num += 1;
497                self.correlated_indices_collection.push(
498                    subquery.collect_correlated_indices_by_scope_and_assign_id(
499                        subquery.outer_scope,
500                        self.correlated_id.unwrap(),
501                    ),
502                );
503                self.subqueries.push(subquery);
504                input_ref
505            }
506        }
507
508        let mut rewriter = SubstituteSubQueries {
509            input_col_num: root.schema().len(),
510            subqueries: vec![],
511            correlated_id: None,
512            correlated_indices_collection: vec![],
513            ctx: self.ctx.clone(),
514        };
515        exprs = exprs
516            .into_iter()
517            .map(|e| rewriter.rewrite_expr(e))
518            .collect();
519
520        let mut right = None;
521
522        for subquery in rewriter.subqueries {
523            let return_type = subquery.return_type();
524            let subroot = self.plan_query(subquery.query)?;
525
526            let subplan = match subquery.kind {
527                SubqueryKind::Scalar => subroot.into_unordered_subplan(),
528                SubqueryKind::UpdateSet => {
529                    let plan = subroot.into_unordered_subplan();
530
531                    // Compose all input columns into a struct with `ROW` function.
532                    let all_input_refs = plan
533                        .schema()
534                        .data_types()
535                        .into_iter()
536                        .enumerate()
537                        .map(|(i, data_type)| InputRef::new(i, data_type).into())
538                        .collect::<Vec<_>>();
539                    let call =
540                        FunctionCall::new_unchecked(ExprType::Row, all_input_refs, return_type);
541
542                    LogicalProject::create(plan, vec![call.into()])
543                }
544                SubqueryKind::Existential => {
545                    self.create_exists(subroot.into_unordered_subplan())?
546                }
547                SubqueryKind::Array => subroot.into_array_agg()?,
548                _ => bail_not_implemented!(issue = 1343, "{:?}", subquery.kind),
549            };
550            if right.is_none() {
551                right = Some(subplan);
552            } else {
553                right = Some(LogicalJoin::create(
554                    right.clone().unwrap(),
555                    subplan,
556                    JoinType::FullOuter,
557                    ExprImpl::literal_bool(true),
558                ));
559            }
560        }
561
562        root = if let Some(right) = right {
563            let mut correlated_indices = rewriter
564                .correlated_indices_collection
565                .iter()
566                .flatten()
567                .cloned()
568                .collect::<Vec<_>>();
569            correlated_indices.sort();
570            correlated_indices.dedup();
571
572            Self::create_apply(
573                rewriter.correlated_id.expect("must have a correlated id"),
574                correlated_indices,
575                root,
576                right,
577                ExprImpl::literal_bool(true),
578                JoinType::LeftOuter,
579                true,
580            )
581        } else {
582            root
583        };
584
585        Ok((root, exprs))
586    }
587
588    fn create_apply(
589        correlated_id: CorrelatedId,
590        correlated_indices: Vec<usize>,
591        left: PlanRef,
592        right: PlanRef,
593        on: ExprImpl,
594        join_type: JoinType,
595        max_one_row: bool,
596    ) -> PlanRef {
597        LogicalApply::create(
598            left,
599            right,
600            join_type,
601            Condition::with_expr(on),
602            correlated_id,
603            correlated_indices,
604            max_one_row,
605        )
606    }
607}
risingwave_frontend/planner/select.rs

risingwave_frontend/planner/
select.rs
