risingwave_frontend/optimizer/plan_visitor/
cardinality_visitor.rs

1// Copyright 2023 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::HashSet;
16use std::ops::{Mul, Sub};
17
18use itertools::Itertools;
19use risingwave_common::types::DataType;
20use risingwave_common::util::iter_util::ZipEqFast;
21use risingwave_pb::plan_common::JoinType;
22
23use super::{DefaultBehavior, DefaultValue, LogicalPlanVisitor, PlanVisitor};
24use crate::expr::{Expr, ExprImpl, ExprType};
25use crate::optimizer::plan_node::generic::TopNLimit;
26use crate::optimizer::plan_node::{
27    self, LogicalPlanRef as PlanRef, PlanTreeNode, PlanTreeNodeBinary, PlanTreeNodeUnary,
28};
29use crate::optimizer::plan_visitor::LogicalPlanNode;
30use crate::optimizer::property::Cardinality;
31use crate::utils::Condition;
32
33/// A visitor that computes the cardinality of a plan node.
34pub struct CardinalityVisitor;
35
36impl CardinalityVisitor {
37    /// Used for `Filter` and `Scan` with predicate.
38    fn visit_predicate(
39        input: &dyn LogicalPlanNode,
40        input_card: Cardinality,
41        eq_set: HashSet<usize>,
42    ) -> Cardinality {
43        // TODO: there could be more unique keys than the stream key after we support it.
44        let unique_keys: Vec<HashSet<_>> = input
45            .stream_key()
46            .into_iter()
47            .map(|s| s.iter().copied().collect())
48            .collect();
49
50        if unique_keys
51            .iter()
52            .any(|unique_key| eq_set.is_superset(unique_key))
53        {
54            input_card.min(0..=1)
55        } else {
56            input_card.min(0..)
57        }
58    }
59
60    /// Returns the columns that `predicate` fixes to a single value:
61    /// - columns compared by `=`, `IS NOT DISTINCT FROM`, or a single-item `IN` with a pure
62    ///   expression that references no input column, such as a constant, a parameter, or a
63    ///   correlated input ref. A row comparison such as `(a, b) = (1, 2)` fixes its columns field
64    ///   by field.
65    /// - boolean columns used as a conjunction directly or through `NOT`, `IS TRUE`, or
66    ///   `IS FALSE`, such as `WHERE flag`.
67    ///
68    /// Other conjunctions, such as `pk > 1`, `pk = other_col`, or `pk = 1 OR pk = 2`, can keep
69    /// several rows even when they reference every unique key column, so they are ignored.
70    fn fixed_columns(predicate: &Condition) -> HashSet<usize> {
71        let mut fixed = HashSet::new();
72        for conjunction in &predicate.conjunctions {
73            match conjunction {
74                ExprImpl::InputRef(input_ref) => {
75                    fixed.insert(input_ref.index());
76                }
77                ExprImpl::FunctionCall(call) => match (call.func_type(), call.inputs()) {
78                    (
79                        ExprType::Not | ExprType::IsTrue | ExprType::IsFalse,
80                        [ExprImpl::InputRef(input_ref)],
81                    ) => {
82                        fixed.insert(input_ref.index());
83                    }
84                    // `In` with two inputs is `lhs IN (rhs)`.
85                    (ExprType::Equal | ExprType::IsNotDistinctFrom | ExprType::In, [lhs, rhs]) => {
86                        Self::collect_equality_fixed_columns(lhs, rhs, &mut fixed);
87                    }
88                    _ => {}
89                },
90                _ => {}
91            }
92        }
93        fixed
94    }
95
96    /// Collects the columns that `lhs = rhs` fixes. When both sides are rows, each pair of
97    /// fields is compared on its own.
98    fn collect_equality_fixed_columns(lhs: &ExprImpl, rhs: &ExprImpl, fixed: &mut HashSet<usize>) {
99        if let (Some(lhs_fields), Some(rhs_fields)) = (Self::row_fields(lhs), Self::row_fields(rhs))
100            && lhs_fields.len() == rhs_fields.len()
101        {
102            for (lhs, rhs) in lhs_fields.iter().zip_eq_fast(rhs_fields) {
103                Self::collect_equality_fixed_columns(lhs, rhs, fixed);
104            }
105            return;
106        }
107        let is_fixed_value = |e: &ExprImpl| !e.has_input_ref() && !e.has_subquery() && e.is_pure();
108        if is_fixed_value(rhs) {
109            Self::collect_operand_columns(lhs, fixed);
110        } else if is_fixed_value(lhs) {
111            Self::collect_operand_columns(rhs, fixed);
112        }
113    }
114
115    /// Collects the columns of an operand that is equal to a fixed value: the column itself, a
116    /// lossless cast of it, or each such field of a row, as in `ROW(a, b) = '(1,2)'`.
117    fn collect_operand_columns(expr: &ExprImpl, fixed: &mut HashSet<usize>) {
118        if let Some(fields) = Self::row_fields(expr) {
119            for field in fields {
120                Self::collect_operand_columns(field, fixed);
121            }
122        } else if let Some(index) = Self::as_injective_input_ref(expr) {
123            fixed.insert(index);
124        }
125    }
126
127    fn row_fields(expr: &ExprImpl) -> Option<&[ExprImpl]> {
128        match expr {
129            ExprImpl::FunctionCall(call) if call.func_type() == ExprType::Row => {
130                Some(call.inputs())
131            }
132            _ => None,
133        }
134    }
135
136    /// Returns the column referenced by `expr` if `expr` is the column itself or a lossless cast
137    /// of it, so that equal results imply equal column values. Such casts come from comparing
138    /// columns of different numeric types, e.g. `CAST(int_pk AS DECIMAL) = 1.5` or
139    /// `CAST(int_pk AS BIGINT) = bigint_col`.
140    fn as_injective_input_ref(expr: &ExprImpl) -> Option<usize> {
141        match expr {
142            ExprImpl::InputRef(input_ref) => Some(input_ref.index()),
143            ExprImpl::FunctionCall(call) if call.func_type() == ExprType::Cast => {
144                let [input] = call.inputs() else {
145                    return None;
146                };
147                let lossless = matches!(
148                    (input.return_type(), call.return_type()),
149                    (
150                        DataType::Int16,
151                        DataType::Int32
152                            | DataType::Int64
153                            | DataType::Decimal
154                            | DataType::Float32
155                            | DataType::Float64
156                    ) | (
157                        DataType::Int32,
158                        DataType::Int64 | DataType::Decimal | DataType::Float64
159                    ) | (DataType::Int64, DataType::Decimal)
160                );
161                if lossless {
162                    Self::as_injective_input_ref(input)
163                } else {
164                    None
165                }
166            }
167            _ => None,
168        }
169    }
170}
171
172impl LogicalPlanVisitor for CardinalityVisitor {
173    type Result = Cardinality;
174
175    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
176
177    fn default_behavior() -> Self::DefaultBehavior {
178        // returns unknown cardinality for default behavior, which is always correct
179        DefaultValue
180    }
181
182    fn visit_logical_values(&mut self, plan: &plan_node::LogicalValues) -> Cardinality {
183        plan.rows().len().into()
184    }
185
186    fn visit_logical_share(&mut self, plan: &plan_node::LogicalShare) -> Cardinality {
187        self.visit(plan.input())
188    }
189
190    fn visit_logical_dedup(&mut self, plan: &plan_node::LogicalDedup) -> Cardinality {
191        let input = self.visit(plan.input());
192        if plan.dedup_cols().is_empty() {
193            input.min(1)
194        } else {
195            input
196        }
197    }
198
199    fn visit_logical_over_window(&mut self, plan: &super::LogicalOverWindow) -> Self::Result {
200        self.visit(plan.input())
201    }
202
203    fn visit_logical_agg(&mut self, plan: &plan_node::LogicalAgg) -> Cardinality {
204        let input = self.visit(plan.input());
205
206        if plan.group_key().is_empty() {
207            // A simple agg always yields exactly one row, even over an empty input.
208            1.into()
209        } else {
210            input.min(1..)
211        }
212    }
213
214    fn visit_logical_limit(&mut self, plan: &plan_node::LogicalLimit) -> Cardinality {
215        self.visit(plan.input()).min(plan.limit() as usize)
216    }
217
218    fn visit_logical_max_one_row(&mut self, plan: &plan_node::LogicalMaxOneRow) -> Cardinality {
219        self.visit(plan.input()).min(1)
220    }
221
222    fn visit_logical_project(&mut self, plan: &plan_node::LogicalProject) -> Cardinality {
223        self.visit(plan.input())
224    }
225
226    fn visit_logical_top_n(&mut self, plan: &plan_node::LogicalTopN) -> Cardinality {
227        let input = self.visit(plan.input());
228
229        let each_group = match plan.limit_attr() {
230            TopNLimit::Simple(limit) => input.sub(plan.offset() as usize).min(limit as usize),
231            TopNLimit::WithTies(limit) => {
232                assert_eq!(plan.offset(), 0, "ties with offset is not supported yet");
233                input.min((limit as usize)..)
234            }
235        };
236
237        if plan.group_key().is_empty() {
238            each_group
239        } else {
240            let group_number = input.min(1..);
241            each_group
242                .mul(group_number)
243                // the output cardinality will never be more than the input, thus `.min(input)`
244                .min(input)
245        }
246    }
247
248    fn visit_logical_filter(&mut self, plan: &plan_node::LogicalFilter) -> Cardinality {
249        let eq_set = Self::fixed_columns(plan.predicate());
250        Self::visit_predicate(&*plan.input(), self.visit(plan.input()), eq_set)
251    }
252
253    fn visit_logical_scan(&mut self, plan: &plan_node::LogicalScan) -> Cardinality {
254        // The predicate refers to table columns, while the stream key refers to output columns.
255        let fixed_table_columns = Self::fixed_columns(plan.predicate());
256        let eq_set = plan
257            .output_col_idx()
258            .iter()
259            .positions(|table_column| fixed_table_columns.contains(table_column))
260            .collect();
261        Self::visit_predicate(plan, plan.table_cardinality(), eq_set)
262    }
263
264    fn visit_logical_union(&mut self, plan: &plan_node::LogicalUnion) -> Cardinality {
265        let all = plan
266            .inputs()
267            .into_iter()
268            .map(|input| self.visit(input))
269            .fold(Cardinality::unknown(), std::ops::Add::add);
270
271        if plan.all() { all } else { all.min(1..) }
272    }
273
274    fn visit_logical_join(&mut self, plan: &plan_node::LogicalJoin) -> Cardinality {
275        let left = self.visit(plan.left());
276        let right = self.visit(plan.right());
277        join_cardinality(plan.join_type(), left, right)
278    }
279
280    fn visit_logical_apply(&mut self, plan: &plan_node::LogicalApply) -> Cardinality {
281        // The right side is evaluated once per left row, so its bound composes like a join.
282        let left = self.visit(plan.left());
283        let right = self.visit(plan.right());
284        join_cardinality(plan.join_type(), left, right)
285    }
286
287    fn visit_logical_now(&mut self, plan: &plan_node::LogicalNow) -> Cardinality {
288        if plan.max_one_row() {
289            1.into()
290        } else {
291            Cardinality::unknown()
292        }
293    }
294
295    fn visit_logical_expand(&mut self, plan: &plan_node::LogicalExpand) -> Cardinality {
296        self.visit(plan.input()) * plan.column_subsets().len()
297    }
298}
299
300fn join_cardinality(join_type: JoinType, left: Cardinality, right: Cardinality) -> Cardinality {
301    match join_type {
302        JoinType::Unspecified => unreachable!(),
303
304        // For each row from one side, we match `0..=(right.hi)` rows from the other side.
305        JoinType::Inner => left.mul(right.min(0..)),
306
307        // For each row from one side, we match `1..=max(right.hi, 1)` rows from the other side,
308        // since we can at least match a `NULL` row.
309        JoinType::LeftOuter => left.mul(right.max(1).min(1..)),
310        JoinType::RightOuter => right.mul(left.max(1).min(1..)),
311
312        // For each row in the result set, it must belong to the given side.
313        JoinType::LeftSemi | JoinType::LeftAnti => left.min(0..),
314        JoinType::RightSemi | JoinType::RightAnti => right.min(0..),
315
316        // TODO: refine the cardinality of full outer join
317        JoinType::FullOuter => {
318            if left.is_at_most(1) && right.is_at_most(1) {
319                Cardinality::new(0, 1)
320            } else {
321                Cardinality::unknown()
322            }
323        }
324
325        // For each row from one side, we match `0..=1` rows from the other side.
326        JoinType::AsofInner => left.mul(right.min(0..=1)),
327        // For each row from left side, we match exactly 1 row from the right side or a `NULL` row`.
328        JoinType::AsofLeftOuter => left,
329    }
330}
331
332#[easy_ext::ext(LogicalCardinalityExt)]
333pub impl PlanRef {
334    /// Returns `true` if the plan node is guaranteed to yield at most one row.
335    fn max_one_row(&self) -> bool {
336        CardinalityVisitor.visit(self.clone()).is_at_most(1)
337    }
338
339    /// Returns the number of rows the plan node is guaranteed to yield, if known exactly.
340    fn row_count(&self) -> Option<usize> {
341        CardinalityVisitor.visit(self.clone()).get_exact()
342    }
343}
risingwave_frontend/optimizer/plan_visitor/cardinality_visitor.rs

risingwave_frontend/optimizer/plan_visitor/
cardinality_visitor.rs
