risingwave_frontend/optimizer/plan_node/
logical_gap_fill.rs

1// Copyright 2025 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_expr::bail;
18
19use super::stream::StreamPlanNodeMetadata;
20use super::{
21    ColPrunable, ColumnPruningContext, ExprRewritable, ExprVisitable, Logical,
22    LogicalPlanRef as PlanRef, LogicalProject, PlanBase, PlanTreeNodeUnary, PredicatePushdown,
23    PredicatePushdownContext, ToBatch, ToStream, ToStreamContext, gen_filter_and_pushdown, generic,
24};
25use crate::binder::BoundFillStrategy;
26use crate::error::Result;
27use crate::expr::{ExprImpl, InputRef};
28use crate::optimizer::plan_node::utils::impl_distill_by_unit;
29use crate::utils::{ColIndexMapping, Condition};
30
31/// `LogicalGapFill` implements [`super::Logical`] to represent a gap-filling operation on a time
32/// series.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct LogicalGapFill {
35    pub base: PlanBase<super::Logical>,
36    core: generic::GapFill<PlanRef>,
37}
38
39impl LogicalGapFill {
40    pub fn new(
41        input: PlanRef,
42        time_col: InputRef,
43        interval: ExprImpl,
44        fill_strategies: Vec<BoundFillStrategy>,
45        partition_by_cols: Vec<InputRef>,
46    ) -> Self {
47        let core = generic::GapFill {
48            input,
49            time_col,
50            interval,
51            fill_strategies,
52            partition_by_cols,
53        };
54        let base = PlanBase::new_logical_with_core(&core);
55        Self { base, core }
56    }
57
58    pub fn time_col(&self) -> &InputRef {
59        &self.core.time_col
60    }
61
62    pub fn interval(&self) -> &ExprImpl {
63        &self.core.interval
64    }
65
66    pub fn fill_strategies(&self) -> &[BoundFillStrategy] {
67        &self.core.fill_strategies
68    }
69
70    pub fn partition_by_cols(&self) -> &[InputRef] {
71        &self.core.partition_by_cols
72    }
73}
74
75impl PlanTreeNodeUnary<Logical> for LogicalGapFill {
76    fn input(&self) -> PlanRef {
77        self.core.input.clone()
78    }
79
80    fn clone_with_input(&self, input: PlanRef) -> Self {
81        Self::new(
82            input,
83            self.time_col().clone(),
84            self.interval().clone(),
85            self.fill_strategies().to_vec(),
86            self.partition_by_cols().to_vec(),
87        )
88    }
89}
90
91impl_plan_tree_node_for_unary! { Logical, LogicalGapFill }
92impl_distill_by_unit!(LogicalGapFill, core, "LogicalGapFill");
93
94impl ColPrunable for LogicalGapFill {
95    fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef {
96        let input_col_num = self.input().schema().len();
97        let mut input_required = FixedBitSet::with_capacity(input_col_num);
98        for &required_col in required_cols {
99            input_required.insert(required_col);
100        }
101        input_required.insert(self.time_col().index());
102        input_required.union_with(&self.interval().collect_input_refs(input_col_num));
103        self.fill_strategies()
104            .iter()
105            .for_each(|s| input_required.insert(s.target_col.index()));
106        self.partition_by_cols()
107            .iter()
108            .for_each(|c| input_required.insert(c.index()));
109
110        let input_required_cols: Vec<_> = input_required.ones().collect();
111        let mut col_index_mapping =
112            ColIndexMapping::with_remaining_columns(&input_required_cols, input_col_num);
113
114        let mut new_core = self.core.clone();
115        new_core.input = self.input().prune_col(&input_required_cols, ctx);
116        new_core.rewrite_with_col_index_mapping(&mut col_index_mapping);
117
118        let logical_gap_fill = Self {
119            base: PlanBase::new_logical_with_core(&new_core),
120            core: new_core,
121        }
122        .into();
123
124        if input_required_cols == required_cols {
125            logical_gap_fill
126        } else {
127            let output_required_cols = required_cols
128                .iter()
129                .map(|&idx| col_index_mapping.map(idx))
130                .collect_vec();
131            let src_size = logical_gap_fill.schema().len();
132            LogicalProject::with_mapping(
133                logical_gap_fill,
134                ColIndexMapping::with_remaining_columns(&output_required_cols, src_size),
135            )
136            .into()
137        }
138    }
139}
140
141impl ExprRewritable<Logical> for LogicalGapFill {
142    fn has_rewritable_expr(&self) -> bool {
143        true
144    }
145
146    fn rewrite_exprs(&self, r: &mut dyn crate::expr::ExprRewriter) -> PlanRef {
147        let mut core = self.core.clone();
148        core.rewrite_exprs(r);
149        Self {
150            base: self.base.clone(),
151            core,
152        }
153        .into()
154    }
155}
156
157impl ExprVisitable for LogicalGapFill {
158    fn visit_exprs(&self, v: &mut dyn crate::expr::ExprVisitor) {
159        self.core.visit_exprs(v);
160    }
161}
162
163impl PredicatePushdown for LogicalGapFill {
164    fn predicate_pushdown(
165        &self,
166        predicate: Condition,
167        ctx: &mut PredicatePushdownContext,
168    ) -> PlanRef {
169        // No pushdown, but keep recursing with a `true` condition so that shares below
170        // still receive a contribution from this parent (see `PredicatePushdownContext`).
171        gen_filter_and_pushdown(self, predicate, Condition::true_cond(), ctx)
172    }
173}
174
175impl ToBatch for LogicalGapFill {
176    fn to_batch(&self) -> Result<super::BatchPlanRef> {
177        bail!("BatchGapFill is not implemented yet")
178    }
179}
180
181impl ToStream for LogicalGapFill {
182    fn to_stream(&self, ctx: &mut ToStreamContext) -> Result<super::StreamPlanRef> {
183        use super::{StreamEowcGapFill, StreamEowcSort, StreamGapFill};
184        use crate::error::ErrorCode;
185        use crate::optimizer::property::RequiredDist;
186
187        let stream_input = self.input().to_stream(ctx)?;
188        let partition_cols = &self.core.partition_by_cols;
189
190        // Choose distribution based on partition_by_cols
191        let new_input = if partition_cols.is_empty() {
192            // No partition: singleton distribution (backward compatible)
193            RequiredDist::single().streaming_enforce_if_not_satisfies(stream_input)?
194        } else {
195            // With partition: hash distribute by partition columns
196            let partition_indices: Vec<usize> = partition_cols.iter().map(|c| c.index()).collect();
197            RequiredDist::shard_by_key(stream_input.schema().len(), &partition_indices)
198                .streaming_enforce_if_not_satisfies(stream_input)?
199        };
200
201        // No stream key validation - gap fill manages its own state internally
202        // using (partition_cols, time_col, upstream_stream_key) as state table PK.
203
204        if ctx.emit_on_window_close() {
205            // EOWC mode requires watermark on time column
206            let time_col_idx = self.core.time_col.index();
207            let input_watermark_cols = new_input.watermark_columns();
208            if !input_watermark_cols.contains(time_col_idx) {
209                return Err(ErrorCode::NotSupported(
210                    "GAP_FILL with EMIT ON WINDOW CLOSE requires a watermark on the time column."
211                        .to_owned(),
212                    format!(
213                        "Please define a watermark on the time column (column index {}).",
214                        time_col_idx
215                    ),
216                )
217                .into());
218            }
219
220            let sorted_input = StreamEowcSort::new(new_input, time_col_idx).into();
221            let core = generic::GapFill {
222                input: sorted_input,
223                time_col: self.core.time_col.clone(),
224                interval: self.core.interval.clone(),
225                fill_strategies: self.core.fill_strategies.clone(),
226                partition_by_cols: self.core.partition_by_cols.clone(),
227            };
228            Ok(StreamEowcGapFill::new(core).into())
229        } else {
230            let core = generic::GapFill {
231                input: new_input,
232                time_col: self.core.time_col.clone(),
233                interval: self.core.interval.clone(),
234                fill_strategies: self.core.fill_strategies.clone(),
235                partition_by_cols: self.core.partition_by_cols.clone(),
236            };
237            Ok(StreamGapFill::new(core).into())
238        }
239    }
240
241    fn logical_rewrite_for_stream(
242        &self,
243        _ctx: &mut super::convert::RewriteStreamContext,
244    ) -> Result<(PlanRef, ColIndexMapping)> {
245        let (input, mut col_index_mapping) = self.input().logical_rewrite_for_stream(_ctx)?;
246        let mut new_core = self.core.clone();
247        new_core.input = input;
248
249        if col_index_mapping.is_identity() {
250            return Ok((
251                Self {
252                    base: PlanBase::new_logical_with_core(&new_core),
253                    core: new_core,
254                }
255                .into(),
256                col_index_mapping,
257            ));
258        }
259
260        new_core.rewrite_with_col_index_mapping(&mut col_index_mapping);
261
262        Ok((
263            Self {
264                base: PlanBase::new_logical_with_core(&new_core),
265                core: new_core,
266            }
267            .into(),
268            col_index_mapping,
269        ))
270    }
271}
risingwave_frontend/optimizer/plan_node/logical_gap_fill.rs

risingwave_frontend/optimizer/plan_node/
logical_gap_fill.rs
