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            // `Array` subqueries decorrelate to a scalar aggregate with
442            // `Coalesce(array_agg, ARRAY[])`, guaranteeing one row per outer key, so
443            // `Inner` is equivalent to `LeftOuter` in batch. `Inner` also avoids the
444            // streaming transient-NULL caused by `StreamHashJoin(LeftOuter)`
445            // null-padding the outer row before the right-side aggregate arrives.
446            let join_type = if matches!(subquery.kind, SubqueryKind::Array) {
447                JoinType::Inner
448            } else {
449                JoinType::LeftOuter
450            };
451            root = Self::create_apply(
452                correlated_id,
453                correlated_indices,
454                root,
455                right,
456                ExprImpl::literal_bool(true),
457                join_type,
458                true,
459            );
460        }
461        Ok((root, exprs))
462    }
463
464    /// Substitutes all [`Subquery`] in `exprs` in a cross join way.
465    ///
466    /// Each time a [`Subquery`] is found, it is replaced by a new [`InputRef`]. And `root` is
467    /// replaced by a new `LeftOuter` [`LogicalApply`] whose left side is `root` and right side is
468    /// the planned subquery.
469    ///
470    /// The [`InputRef`]s' indexes start from `root.schema().len()`,
471    /// which means they are additional columns beyond the original `root`.
472    ///
473    ///
474    /// The cross join way only meaningful when there are multiple subqueries
475    ///
476    /// ```text
477    ///            Apply
478    ///           /    \
479    ///         left   CrossJoin
480    ///                   /   \
481    ///           Subquery1   CrossJoin
482    ///                         /   \
483    ///                 Subquery2   Subquery3
484    /// ```
485    /// Typically, this is used for scalar subqueries in select clauses, because users might write subqueries in exprs multiple times.
486    /// If we use the left-deep tree way, it will generate a lot of `Apply` nodes, which is not efficient.
487    pub(super) fn substitute_subqueries_in_cross_join_way(
488        &mut self,
489        mut root: PlanRef,
490        mut exprs: Vec<ExprImpl>,
491    ) -> Result<(PlanRef, Vec<ExprImpl>)> {
492        struct SubstituteSubQueries {
493            input_col_num: usize,
494            subqueries: Vec<Subquery>,
495            correlated_id: Option<CorrelatedId>,
496            correlated_indices_collection: Vec<Vec<usize>>,
497            ctx: OptimizerContextRef,
498        }
499
500        impl ExprRewriter for SubstituteSubQueries {
501            fn rewrite_subquery(&mut self, mut subquery: Subquery) -> ExprImpl {
502                if self.correlated_id.is_none() {
503                    self.correlated_id = Some(self.ctx.next_correlated_id());
504                }
505                let input_ref = InputRef::new(self.input_col_num, subquery.return_type()).into();
506                self.input_col_num += 1;
507                self.correlated_indices_collection.push(
508                    subquery.collect_correlated_indices_by_scope_and_assign_id(
509                        subquery.outer_scope,
510                        self.correlated_id.unwrap(),
511                    ),
512                );
513                self.subqueries.push(subquery);
514                input_ref
515            }
516        }
517
518        let mut rewriter = SubstituteSubQueries {
519            input_col_num: root.schema().len(),
520            subqueries: vec![],
521            correlated_id: None,
522            correlated_indices_collection: vec![],
523            ctx: self.ctx.clone(),
524        };
525        exprs = exprs
526            .into_iter()
527            .map(|e| rewriter.rewrite_expr(e))
528            .collect();
529
530        let mut right = None;
531        let mut all_array_subqueries = true;
532
533        for subquery in rewriter.subqueries {
534            let return_type = subquery.return_type();
535            let subroot = self.plan_query(subquery.query)?;
536
537            let subplan = match subquery.kind {
538                SubqueryKind::Scalar => subroot.into_unordered_subplan(),
539                SubqueryKind::UpdateSet => {
540                    let plan = subroot.into_unordered_subplan();
541
542                    // Compose all input columns into a struct with `ROW` function.
543                    let all_input_refs = plan
544                        .schema()
545                        .data_types()
546                        .into_iter()
547                        .enumerate()
548                        .map(|(i, data_type)| InputRef::new(i, data_type).into())
549                        .collect::<Vec<_>>();
550                    let call =
551                        FunctionCall::new_unchecked(ExprType::Row, all_input_refs, return_type);
552
553                    LogicalProject::create(plan, vec![call.into()])
554                }
555                SubqueryKind::Existential => {
556                    self.create_exists(subroot.into_unordered_subplan())?
557                }
558                SubqueryKind::Array => subroot.into_array_agg()?,
559                _ => bail_not_implemented!(issue = 1343, "{:?}", subquery.kind),
560            };
561            if !matches!(subquery.kind, SubqueryKind::Array) {
562                all_array_subqueries = false;
563            }
564            if right.is_none() {
565                right = Some(subplan);
566            } else {
567                right = Some(LogicalJoin::create(
568                    right.clone().unwrap(),
569                    subplan,
570                    JoinType::FullOuter,
571                    ExprImpl::literal_bool(true),
572                ));
573            }
574        }
575
576        root = if let Some(right) = right {
577            let mut correlated_indices = rewriter
578                .correlated_indices_collection
579                .iter()
580                .flatten()
581                .cloned()
582                .collect::<Vec<_>>();
583            correlated_indices.sort();
584            correlated_indices.dedup();
585
586            // Same reasoning as `substitute_subqueries_in_left_deep_tree_way`:
587            // only safe when every combined subquery is `Array` (Scalar/Existential
588            // legitimately need `LeftOuter` null-padding).
589            let join_type = if all_array_subqueries {
590                JoinType::Inner
591            } else {
592                JoinType::LeftOuter
593            };
594            Self::create_apply(
595                rewriter.correlated_id.expect("must have a correlated id"),
596                correlated_indices,
597                root,
598                right,
599                ExprImpl::literal_bool(true),
600                join_type,
601                true,
602            )
603        } else {
604            root
605        };
606
607        Ok((root, exprs))
608    }
609
610    fn create_apply(
611        correlated_id: CorrelatedId,
612        correlated_indices: Vec<usize>,
613        left: PlanRef,
614        right: PlanRef,
615        on: ExprImpl,
616        join_type: JoinType,
617        max_one_row: bool,
618    ) -> PlanRef {
619        LogicalApply::create(
620            left,
621            right,
622            join_type,
623            Condition::with_expr(on),
624            correlated_id,
625            correlated_indices,
626            max_one_row,
627        )
628    }
629}
risingwave_frontend/planner/select.rs

risingwave_frontend/planner/
select.rs
