risingwave_frontend/optimizer/plan_node/
logical_dedup.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 fixedbitset::FixedBitSet;
16use itertools::Itertools;
17use risingwave_common::util::column_index_mapping::ColIndexMapping;
18
19use super::generic::{GenericPlanRef, TopNLimit};
20use super::utils::impl_distill_by_unit;
21use super::{
22    BatchGroupTopN, BatchPlanRef, ColPrunable, ColumnPruningContext, ExprRewritable, Logical,
23    LogicalPlanRef as PlanRef, LogicalProject, PlanBase, PlanTreeNodeUnary, PredicatePushdown,
24    PredicatePushdownContext, RewriteStreamContext, StreamDedup, StreamGroupTopN, ToBatch,
25    ToStream, ToStreamContext, gen_filter_and_pushdown, generic, try_enforce_locality_requirement,
26};
27use crate::error::Result;
28use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
29use crate::optimizer::property::{Order, RequiredDist};
30use crate::utils::Condition;
31
32/// [`LogicalDedup`] deduplicates data on specific columns. It is now used in `DISTINCT ON` without
33/// an `ORDER BY`.
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub struct LogicalDedup {
36    pub base: PlanBase<Logical>,
37    core: generic::Dedup<PlanRef>,
38}
39
40impl LogicalDedup {
41    pub fn new(input: PlanRef, dedup_cols: Vec<usize>) -> Self {
42        let core = generic::Dedup::new(input, dedup_cols);
43        let base = PlanBase::new_logical_with_core(&core);
44        LogicalDedup { base, core }
45    }
46
47    pub fn dedup_cols(&self) -> &[usize] {
48        &self.core.dedup_cols
49    }
50}
51
52impl PlanTreeNodeUnary<Logical> for LogicalDedup {
53    fn input(&self) -> PlanRef {
54        self.core.input.clone()
55    }
56
57    fn clone_with_input(&self, input: PlanRef) -> Self {
58        Self::new(input, self.dedup_cols().to_vec())
59    }
60
61    fn rewrite_with_input(
62        &self,
63        input: PlanRef,
64        input_col_change: ColIndexMapping,
65    ) -> (Self, ColIndexMapping) {
66        (
67            Self::new(
68                input,
69                self.dedup_cols()
70                    .iter()
71                    .map(|idx| input_col_change.map(*idx))
72                    .collect_vec(),
73            ),
74            input_col_change,
75        )
76    }
77}
78
79impl_plan_tree_node_for_unary! { Logical, LogicalDedup}
80
81impl PredicatePushdown for LogicalDedup {
82    fn predicate_pushdown(
83        &self,
84        predicate: Condition,
85        ctx: &mut PredicatePushdownContext,
86    ) -> PlanRef {
87        gen_filter_and_pushdown(self, predicate, Condition::true_cond(), ctx)
88    }
89}
90
91impl ToStream for LogicalDedup {
92    fn try_better_locality(&self, columns: &[usize]) -> Option<PlanRef> {
93        if columns.is_empty() {
94            return None;
95        }
96
97        // Dedup stores rows with dedup columns as the primary key in its internal state
98        // table, so it can directly satisfy locality requests on a prefix of its dedup columns.
99        let dedup_cols = self.dedup_cols();
100        if columns.len() > dedup_cols.len() || columns != &dedup_cols[..columns.len()] {
101            return None;
102        }
103
104        // Similar to agg/topn, return the current plan directly instead of asking input for
105        // better locality first. The locality can be provided by the current Dedup itself
106        // after `to_stream`, while its input does not have it yet during logical rewrite.
107        Some(self.clone_with_input(self.input()).into())
108    }
109
110    fn logical_rewrite_for_stream(
111        &self,
112        ctx: &mut RewriteStreamContext,
113    ) -> Result<(PlanRef, ColIndexMapping)> {
114        let logical_input = try_enforce_locality_requirement(
115            self.input(),
116            self.dedup_cols(),
117            ctx.locality_backfill_enabled(),
118        );
119        let (input, input_col_change) = logical_input.logical_rewrite_for_stream(ctx)?;
120        let (logical, out_col_change) = self.rewrite_with_input(input, input_col_change);
121        Ok((logical.into(), out_col_change))
122    }
123
124    fn to_stream(
125        &self,
126        ctx: &mut ToStreamContext,
127    ) -> Result<crate::optimizer::plan_node::StreamPlanRef> {
128        use super::stream::prelude::*;
129
130        let input = self.input().to_stream(ctx)?;
131        let input = RequiredDist::hash_shard(self.dedup_cols())
132            .streaming_enforce_if_not_satisfies(input)?;
133        if input.append_only() {
134            // `LogicalDedup` is transformed to `StreamDedup` only when the input is append-only.
135            let core = self.core.clone_with_input(input);
136            Ok(StreamDedup::new(core).into())
137        } else {
138            // If the input is not append-only, we use a `StreamGroupTopN` with the limit being 1.
139            let logical_top_n = generic::TopN::with_group(
140                input,
141                TopNLimit::new(1, false),
142                0,
143                Order::default(),
144                self.dedup_cols().to_vec(),
145            );
146            Ok(StreamGroupTopN::new(logical_top_n, None)?.into())
147        }
148    }
149}
150
151impl ToBatch for LogicalDedup {
152    fn to_batch(&self) -> Result<BatchPlanRef> {
153        let input = self.input().to_batch()?;
154        let logical_top_n = generic::TopN::with_group(
155            input,
156            TopNLimit::new(1, false),
157            0,
158            Order::default(),
159            self.dedup_cols().to_vec(),
160        );
161        Ok(BatchGroupTopN::new(logical_top_n).into())
162    }
163}
164
165impl ExprRewritable<Logical> for LogicalDedup {}
166
167impl ExprVisitable for LogicalDedup {}
168
169impl ColPrunable for LogicalDedup {
170    fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef {
171        let input_required_bitset = FixedBitSet::from_iter(required_cols.iter().copied());
172        let dedup_required_bitset = {
173            let mut dedup_required_bitset = FixedBitSet::with_capacity(self.input().schema().len());
174            self.dedup_cols()
175                .iter()
176                .for_each(|idx| dedup_required_bitset.insert(*idx));
177            dedup_required_bitset
178        };
179        let input_required_cols = {
180            let mut tmp = input_required_bitset;
181            tmp.union_with(&dedup_required_bitset);
182            tmp.ones().collect_vec()
183        };
184        let mapping = ColIndexMapping::with_remaining_columns(
185            &input_required_cols,
186            self.input().schema().len(),
187        );
188
189        let new_input = self.input().prune_col(&input_required_cols, ctx);
190        let new_dedup_cols = self
191            .dedup_cols()
192            .iter()
193            .map(|&idx| mapping.map(idx))
194            .collect_vec();
195        let logical_dedup = Self::new(new_input, new_dedup_cols).into();
196
197        if input_required_cols == required_cols {
198            logical_dedup
199        } else {
200            let output_required_cols = required_cols
201                .iter()
202                .map(|&idx| mapping.map(idx))
203                .collect_vec();
204            let src_size = logical_dedup.schema().len();
205            LogicalProject::with_mapping(
206                logical_dedup,
207                ColIndexMapping::with_remaining_columns(&output_required_cols, src_size),
208            )
209            .into()
210        }
211    }
212}
213
214impl_distill_by_unit!(LogicalDedup, core, "LogicalDedup");
risingwave_frontend/optimizer/plan_node/logical_dedup.rs

risingwave_frontend/optimizer/plan_node/
logical_dedup.rs
