risingwave_frontend/optimizer/plan_visitor/
jsonb_stream_key_checker.rs

1// Copyright 2024 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 risingwave_common::catalog::{Field, FieldDisplay};
16use risingwave_common::types::DataType;
17
18use super::{DefaultBehavior, LogicalPlanVisitor, Merge};
19use crate::optimizer::plan_node::generic::GenericPlanRef;
20use crate::optimizer::plan_node::*;
21use crate::optimizer::plan_visitor::PlanVisitor;
22
23/// Finds the first plan position where the checked type would become part of a key, and names the
24/// SQL clause responsible so the error points at what the user wrote.
25///
26/// Positions checked only for `Variant` are deliberately not checked for `Jsonb`: widening the
27/// pre-existing JSONB gate would be a breaking change.
28#[derive(Debug, Clone, Copy)]
29pub enum StreamKeyChecker {
30    Jsonb,
31    Variant,
32}
33
34impl StreamKeyChecker {
35    fn visit_inputs(&mut self, plan: &impl LogicalPlanNode) -> Option<String> {
36        let results = plan.inputs().into_iter().map(|input| self.visit(input));
37        Self::default_behavior().apply(results)
38    }
39
40    fn err_msg(&self, target: &str, field: &Field) -> String {
41        format!(
42            "{} column \"{}\" should not be in the {}.",
43            self.type_name(),
44            FieldDisplay(field),
45            target
46        )
47    }
48
49    fn is_restricted_key_type(&self, data_type: &DataType) -> bool {
50        match self {
51            Self::Jsonb => matches!(data_type, DataType::Jsonb),
52            // Unlike the pre-existing JSONB behavior, VARIANT is rejected even when nested.
53            Self::Variant => data_type.contains_variant(),
54        }
55    }
56
57    fn type_name(&self) -> &'static str {
58        match self {
59            Self::Jsonb => "JSONB",
60            Self::Variant => "VARIANT",
61        }
62    }
63}
64
65impl LogicalPlanVisitor for StreamKeyChecker {
66    type Result = Option<String>;
67
68    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
69
70    fn default_behavior() -> Self::DefaultBehavior {
71        Merge(|a: Option<String>, b| a.or(b))
72    }
73
74    fn visit_logical_dedup(&mut self, plan: &LogicalDedup) -> Self::Result {
75        let input = plan.input();
76        let schema = input.schema();
77        let data_types = schema.data_types();
78        for idx in plan.dedup_cols() {
79            if self.is_restricted_key_type(&data_types[*idx]) {
80                return Some(self.err_msg("distinct key", &schema[*idx]));
81            }
82        }
83        self.visit_inputs(plan)
84    }
85
86    fn visit_logical_top_n(&mut self, plan: &LogicalTopN) -> Self::Result {
87        let input = plan.input();
88        let schema = input.schema();
89        let data_types = schema.data_types();
90        for idx in plan.group_key() {
91            if self.is_restricted_key_type(&data_types[*idx]) {
92                return Some(self.err_msg("TopN group key", &schema[*idx]));
93            }
94        }
95        for idx in plan
96            .topn_order()
97            .column_orders
98            .iter()
99            .map(|c| c.column_index)
100        {
101            if self.is_restricted_key_type(&data_types[idx]) {
102                return Some(self.err_msg("TopN order key", &schema[idx]));
103            }
104        }
105        self.visit_inputs(plan)
106    }
107
108    fn visit_logical_union(&mut self, plan: &LogicalUnion) -> Self::Result {
109        if !plan.all() {
110            for field in &plan.inputs()[0].schema().fields {
111                if self.is_restricted_key_type(&field.data_type()) {
112                    return Some(self.err_msg("field", field));
113                }
114            }
115        }
116        self.visit_inputs(plan)
117    }
118
119    fn visit_logical_agg(&mut self, plan: &LogicalAgg) -> Self::Result {
120        let input = plan.input();
121        let schema = input.schema();
122        let data_types = schema.data_types();
123        for idx in plan.group_key().indices() {
124            if self.is_restricted_key_type(&data_types[idx]) {
125                return Some(self.err_msg("aggregation group key", &schema[idx]));
126            }
127        }
128        // An aggregate call's own ORDER BY / DISTINCT lands in the state table key too.
129        if matches!(self, Self::Variant) {
130            for call in plan.agg_calls() {
131                for idx in call.order_by.iter().map(|c| c.column_index) {
132                    if self.is_restricted_key_type(&data_types[idx]) {
133                        return Some(self.err_msg("aggregation order key", &schema[idx]));
134                    }
135                }
136                for input in call.distinct.then_some(&call.inputs).into_iter().flatten() {
137                    let idx = input.index();
138                    if self.is_restricted_key_type(&data_types[idx]) {
139                        return Some(self.err_msg("distinct aggregation argument", &schema[idx]));
140                    }
141                }
142            }
143        }
144        self.visit_inputs(plan)
145    }
146
147    fn visit_logical_join(&mut self, plan: &LogicalJoin) -> Self::Result {
148        // Only the equi keys become a hash key; the rest of `on` stays a per-row predicate.
149        if matches!(self, Self::Variant) {
150            let left = plan.left();
151            let right = plan.right();
152            let predicate = EqJoinPredicate::create(
153                left.schema().len(),
154                right.schema().len(),
155                plan.on().clone(),
156            );
157            for (left_idx, right_idx) in predicate.eq_indexes() {
158                for (schema, idx) in [(left.schema(), left_idx), (right.schema(), right_idx)] {
159                    if self.is_restricted_key_type(&schema[idx].data_type()) {
160                        return Some(self.err_msg("join key", &schema[idx]));
161                    }
162                }
163            }
164        }
165        self.visit_inputs(plan)
166    }
167
168    fn visit_logical_over_window(&mut self, plan: &LogicalOverWindow) -> Self::Result {
169        let input = plan.input();
170        let schema = input.schema();
171        let data_types = schema.data_types();
172
173        for func in plan.window_functions() {
174            for idx in func.partition_by.iter().map(|e| e.index()) {
175                if self.is_restricted_key_type(&data_types[idx]) {
176                    return Some(self.err_msg("over window partition key", &schema[idx]));
177                }
178            }
179
180            for idx in func.order_by.iter().map(|c| c.column_index) {
181                if self.is_restricted_key_type(&data_types[idx]) {
182                    return Some(self.err_msg("over window order by key", &schema[idx]));
183                }
184            }
185        }
186        self.visit_inputs(plan)
187    }
188}
risingwave_frontend/optimizer/plan_visitor/jsonb_stream_key_checker.rs

risingwave_frontend/optimizer/plan_visitor/
jsonb_stream_key_checker.rs
