risingwave_frontend/optimizer/plan_node/
logical_share.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 pretty_xmlish::{Pretty, XmlNode};
16use risingwave_common::bail_not_implemented;
17
18use super::utils::{Distill, childless_record};
19use super::{
20    ColPrunable, ExprRewritable, Logical, LogicalPlanRef as PlanRef, PlanBase, PlanTreeNodeUnary,
21    PredicatePushdown, ShareNode, StreamPlanRef, ToBatch, ToStream, generic,
22};
23use crate::error::Result;
24use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
25use crate::optimizer::plan_node::generic::{GenericPlanRef, Share};
26use crate::optimizer::plan_node::{
27    ColumnPruningContext, PredicatePushdownContext, RewriteStreamContext, StreamShare,
28    ToStreamContext,
29};
30use crate::optimizer::{OptimizerContextRef, ShareId};
31use crate::utils::{ColIndexMapping, Condition};
32
33/// `LogicalShare` operator is used to represent reusing of existing operators.
34/// It is the key operator for DAG plan.
35/// It could have multiple parents which makes it different from other operators.
36/// Currently, it has been used to the following scenarios:
37/// 1. Share source.
38/// 2. Subquery unnesting domain calculation.
39///
40/// A DAG plan example: A self join shares the same source.
41/// ```text
42///     LogicalJoin
43///    /           \
44///   |            |
45///   \           /
46///   LogicalShare
47///        |
48///   LogicalSource
49/// ```
50#[derive(Debug, Clone, PartialEq, Eq, Hash)]
51pub struct LogicalShare {
52    pub base: PlanBase<Logical>,
53    core: generic::Share<PlanRef>,
54}
55
56impl LogicalShare {
57    pub fn new(input: PlanRef) -> Self {
58        let ctx = input.ctx();
59        let core = ctx.register_logical_share(input);
60        Self::with_core(core)
61    }
62
63    fn with_core(core: generic::Share<PlanRef>) -> Self {
64        let base = PlanBase::new_logical_share(&core);
65        LogicalShare { base, core }
66    }
67
68    pub(in crate::optimizer) fn from_share_id(ctx: OptimizerContextRef, share_id: ShareId) -> Self {
69        Self::with_core(ctx.logical_share(share_id))
70    }
71
72    pub fn create(input: PlanRef) -> PlanRef {
73        LogicalShare::new(input).into()
74    }
75
76    pub(super) fn pretty_fields(base: impl GenericPlanRef, name: &str) -> XmlNode<'_> {
77        childless_record(name, vec![("id", Pretty::debug(&base.id().0))])
78    }
79}
80
81impl PlanTreeNodeUnary<Logical> for LogicalShare {
82    fn input(&self) -> PlanRef {
83        self.core.input()
84    }
85
86    fn clone_with_input(&self, _input: PlanRef) -> Self {
87        unreachable!("shared node should be handled specially in PlanRef::clone_with_input")
88    }
89
90    fn rewrite_with_input(
91        &self,
92        input: PlanRef,
93        input_col_change: ColIndexMapping,
94    ) -> (Self, ColIndexMapping) {
95        (Self::new(input), input_col_change)
96    }
97}
98
99impl_plan_tree_node_for_unary! { Logical, LogicalShare}
100
101impl ShareNode<Logical> for LogicalShare {
102    fn share_id(&self) -> ShareId {
103        self.core.share_id()
104    }
105
106    fn new_share(core: Share<PlanRef>) -> PlanRef {
107        Self::with_core(core).into()
108    }
109
110    fn replace_input(&self, plan: PlanRef) -> PlanRef {
111        debug_assert!(
112            self.schema().type_eq(plan.schema()),
113            "replacing a logical share input must preserve its schema"
114        );
115        self.ctx()
116            .update_logical_share(self.share_id(), plan.clone());
117        Self::with_core(self.core.with_input(plan)).into()
118    }
119
120    fn fork_with_input(&self, plan: PlanRef) -> PlanRef {
121        Self::new(plan).into()
122    }
123}
124
125impl Distill for LogicalShare {
126    fn distill<'a>(&self) -> XmlNode<'a> {
127        Self::pretty_fields(&self.base, "LogicalShare")
128    }
129}
130
131impl ColPrunable for LogicalShare {
132    fn prune_col(&self, _required_cols: &[usize], _ctx: &mut ColumnPruningContext) -> PlanRef {
133        unimplemented!("call prune_col of the PlanRef instead of calling directly on LogicalShare")
134    }
135}
136
137impl ExprRewritable<Logical> for LogicalShare {}
138
139impl ExprVisitable for LogicalShare {}
140
141impl PredicatePushdown for LogicalShare {
142    fn predicate_pushdown(
143        &self,
144        _predicate: Condition,
145        _ctx: &mut PredicatePushdownContext,
146    ) -> PlanRef {
147        unimplemented!(
148            "call predicate_pushdown of the PlanRef instead of calling directly on LogicalShare"
149        )
150    }
151}
152
153impl ToBatch for LogicalShare {
154    fn to_batch(&self) -> Result<crate::optimizer::plan_node::BatchPlanRef> {
155        bail_not_implemented!("batch query doesn't support share operator for now");
156    }
157}
158
159impl ToStream for LogicalShare {
160    fn to_stream(
161        &self,
162        ctx: &mut ToStreamContext,
163    ) -> Result<crate::optimizer::plan_node::StreamPlanRef> {
164        match ctx.get_to_stream_result(self.share_id()) {
165            None => {
166                let new_input = self.input().to_stream(ctx)?;
167                let stream_share_ref: StreamPlanRef = StreamShare::new_from_input(new_input).into();
168                ctx.add_to_stream_result(self.share_id(), stream_share_ref.clone());
169                Ok(stream_share_ref)
170            }
171            Some(cache) => Ok(cache.clone()),
172        }
173    }
174
175    fn logical_rewrite_for_stream(
176        &self,
177        ctx: &mut RewriteStreamContext,
178    ) -> Result<(PlanRef, ColIndexMapping)> {
179        match ctx.get_rewrite_result(self.share_id()) {
180            None => {
181                let (new_input, col_change) = self.input().logical_rewrite_for_stream(ctx)?;
182                let new_share: PlanRef = Self::new(new_input).into();
183                ctx.add_rewrite_result(self.share_id(), new_share.clone(), col_change.clone());
184                Ok((new_share, col_change))
185            }
186            Some(cache) => Ok(cache.clone()),
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193
194    use risingwave_common::catalog::{Field, Schema};
195    use risingwave_common::types::{DataType, ScalarImpl};
196    use risingwave_pb::expr::expr_node::Type;
197    use risingwave_pb::plan_common::JoinType;
198
199    use super::*;
200    use crate::expr::{ExprImpl, FunctionCall, InputRef, Literal};
201    use crate::optimizer::optimizer_context::OptimizerContext;
202    use crate::optimizer::plan_node::{
203        LogicalFilter, LogicalJoin, LogicalValues, PlanTreeNodeBinary,
204    };
205
206    #[tokio::test]
207    async fn test_share_predicate_pushdown() {
208        let ty = DataType::Int32;
209        let ctx = OptimizerContext::mock();
210        let fields: Vec<Field> = vec![
211            Field::with_name(ty.clone(), "v1"),
212            Field::with_name(ty.clone(), "v2"),
213            Field::with_name(ty.clone(), "v3"),
214        ];
215        let values1 = LogicalValues::new(vec![], Schema { fields }, ctx);
216
217        let share: PlanRef = LogicalShare::create(values1.into());
218
219        let on: ExprImpl = ExprImpl::FunctionCall(Box::new(
220            FunctionCall::new(
221                Type::Equal,
222                vec![
223                    ExprImpl::InputRef(Box::new(InputRef::new(1, ty.clone()))),
224                    ExprImpl::InputRef(Box::new(InputRef::new(3, ty))),
225                ],
226            )
227            .unwrap(),
228        ));
229
230        let predicate1: ExprImpl = ExprImpl::FunctionCall(Box::new(
231            FunctionCall::new(
232                Type::Equal,
233                vec![
234                    ExprImpl::InputRef(Box::new(InputRef::new(0, DataType::Int32))),
235                    ExprImpl::Literal(Box::new(Literal::new(
236                        Some(ScalarImpl::from(100)),
237                        DataType::Int32,
238                    ))),
239                ],
240            )
241            .unwrap(),
242        ));
243
244        let predicate2: ExprImpl = ExprImpl::FunctionCall(Box::new(
245            FunctionCall::new(
246                Type::Equal,
247                vec![
248                    ExprImpl::InputRef(Box::new(InputRef::new(4, DataType::Int32))),
249                    ExprImpl::Literal(Box::new(Literal::new(
250                        Some(ScalarImpl::from(200)),
251                        DataType::Int32,
252                    ))),
253                ],
254            )
255            .unwrap(),
256        ));
257
258        let join: PlanRef = LogicalJoin::create(share.clone(), share, JoinType::Inner, on);
259
260        let filter1: PlanRef = LogicalFilter::create_with_expr(join, predicate1);
261
262        let filter2: PlanRef = LogicalFilter::create_with_expr(filter1, predicate2);
263
264        let result = filter2.predicate_pushdown(
265            Condition::true_cond(),
266            &mut PredicatePushdownContext::new(filter2.clone()),
267        );
268
269        // LogicalJoin { type: Inner, on: (v2 = v1) }
270        // ├─LogicalFilter { predicate: (v1 = 100:Int32) }
271        // | └─LogicalShare { id = 2 }
272        // |   └─LogicalFilter { predicate: ((v1 = 100:Int32) OR (v2 = 200:Int32)) }
273        // |     └─LogicalValues { schema: Schema { fields: [v1:Int32, v2:Int32, v3:Int32] } }
274        // └─LogicalFilter { predicate: (v2 = 200:Int32) }
275        //   └─LogicalShare { id = 2 }
276        //     └─LogicalFilter { predicate: ((v1 = 100:Int32) OR (v2 = 200:Int32)) }
277        //       └─LogicalValues { schema: Schema { fields: [v1:Int32, v2:Int32, v3:Int32] } }
278
279        let logical_join: &LogicalJoin = result.as_logical_join().unwrap();
280        let left = logical_join.left();
281        let left_filter: &LogicalFilter = left.as_logical_filter().unwrap();
282        let left_filter_input = left_filter.input();
283        let logical_share: &LogicalShare = left_filter_input.as_logical_share().unwrap();
284        let share_input = logical_share.input();
285        let share_input_filter: &LogicalFilter = share_input.as_logical_filter().unwrap();
286        let disjunctions = share_input_filter.predicate().conjunctions[0]
287            .as_or_disjunctions()
288            .unwrap();
289        assert_eq!(disjunctions.len(), 2);
290        let (input_ref1, _const1) = disjunctions[0].as_eq_const().unwrap();
291        let (input_ref2, _const2) = disjunctions[1].as_eq_const().unwrap();
292        if input_ref1.index() == 0 {
293            assert_eq!(input_ref2.index(), 1);
294        } else {
295            assert_eq!(input_ref1.index(), 1);
296            assert_eq!(input_ref2.index(), 0);
297        }
298    }
299}
risingwave_frontend/optimizer/plan_node/logical_share.rs

risingwave_frontend/optimizer/plan_node/
logical_share.rs
