risingwave_frontend/optimizer/plan_node/
stream_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 risingwave_common::util::sort_util::OrderType;
16use risingwave_pb::stream_plan::stream_node::NodeBody;
17
18use super::generic::GenericPlanNode;
19use super::utils::TableCatalogBuilder;
20use super::{
21    ExprRewritable, ExprVisitable, PlanBase, PlanRef, PlanTreeNodeUnary, Stream, TryToStreamPb,
22    generic,
23};
24use crate::binder::BoundFillStrategy;
25use crate::expr::{Expr, ExprImpl, ExprRewriter, ExprVisitor, InputRef};
26use crate::optimizer::plan_node::stream::StreamPlanNodeMetadata;
27use crate::optimizer::plan_node::utils::impl_distill_by_unit;
28use crate::optimizer::property::{Distribution, MonotonicityMap, WatermarkColumns};
29use crate::scheduler::SchedulerResult;
30use crate::stream_fragmenter::BuildFragmentGraphState;
31
32/// `StreamGapFill` implements [`super::Stream`] to represent a gap-filling operation on a time
33/// series in normal streaming mode (without EOWC semantics).
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub struct StreamGapFill {
36    pub base: PlanBase<super::Stream>,
37    core: generic::GapFill<PlanRef<Stream>>,
38}
39
40impl StreamGapFill {
41    pub fn new(core: generic::GapFill<PlanRef<Stream>>) -> Self {
42        let input = &core.input;
43        let partition_indices = core.partition_key_indices();
44        let distinct_partition_key_count = partition_indices
45            .iter()
46            .copied()
47            .collect::<std::collections::HashSet<_>>()
48            .len();
49        assert_eq!(
50            partition_indices.len(),
51            distinct_partition_key_count,
52            "stream gap fill expects canonicalized partition_by columns",
53        );
54        assert!(
55            !core
56                .pointer_key_indices()
57                .expect("stream gap fill input should have stream key")
58                .is_empty(),
59            "stream gap fill pointer key should not be empty",
60        );
61
62        let dist = if core.partition_by_cols.is_empty() {
63            Distribution::Single
64        } else {
65            Distribution::HashShard(partition_indices)
66        };
67
68        // Gap fill back-fills below the current watermark and emits non-monotonic fill values,
69        // so it preserves neither watermark nor monotonicity on its output.
70        let base = PlanBase::new_stream_with_core(
71            &core,
72            dist,
73            input.stream_kind(),
74            false, // does NOT provide EOWC semantics
75            WatermarkColumns::new(),
76            MonotonicityMap::new(),
77        );
78        Self { base, core }
79    }
80
81    pub fn time_col(&self) -> &InputRef {
82        &self.core.time_col
83    }
84
85    pub fn interval(&self) -> &ExprImpl {
86        &self.core.interval
87    }
88
89    pub fn fill_strategies(&self) -> &[BoundFillStrategy] {
90        &self.core.fill_strategies
91    }
92
93    fn pointer_key_indices(&self) -> Vec<usize> {
94        self.core
95            .pointer_key_indices()
96            .expect("stream gap fill input should have stream key")
97    }
98
99    fn infer_state_table(&self) -> crate::TableCatalog {
100        let mut tbl_builder = TableCatalogBuilder::default();
101
102        let out_schema = self.core.schema();
103        for field in out_schema.fields() {
104            tbl_builder.add_column(field);
105        }
106
107        let state_key_indices = self
108            .core
109            .stream_key_indices()
110            .expect("stream gap fill input should have stream key");
111
112        // PK: deduplicated (partition_cols..., time_col, upstream stream key...).
113        for key_idx in state_key_indices {
114            tbl_builder.add_order_column(key_idx, OrderType::ascending());
115        }
116
117        let partition_key_indices = self.core.partition_key_indices();
118        let read_prefix_len_hint = partition_key_indices.len();
119        tbl_builder.build(partition_key_indices, read_prefix_len_hint)
120    }
121}
122
123impl PlanTreeNodeUnary<Stream> for StreamGapFill {
124    fn input(&self) -> PlanRef<Stream> {
125        self.core.input.clone()
126    }
127
128    fn clone_with_input(&self, input: PlanRef<Stream>) -> Self {
129        let mut core = self.core.clone();
130        core.input = input;
131        Self::new(core)
132    }
133}
134
135impl_plan_tree_node_for_unary! { Stream, StreamGapFill }
136impl_distill_by_unit!(StreamGapFill, core, "StreamGapFill");
137
138impl TryToStreamPb for StreamGapFill {
139    fn try_to_stream_prost_body(
140        &self,
141        state: &mut BuildFragmentGraphState,
142    ) -> SchedulerResult<NodeBody> {
143        use risingwave_pb::stream_plan::*;
144
145        let fill_strategies: Vec<String> = self
146            .fill_strategies()
147            .iter()
148            .map(|strategy| match strategy.strategy {
149                crate::binder::FillStrategy::Locf => "locf".to_owned(),
150                crate::binder::FillStrategy::Interpolate => "interpolate".to_owned(),
151                crate::binder::FillStrategy::Null => "null".to_owned(),
152            })
153            .collect();
154
155        let state_table = self
156            .infer_state_table()
157            .with_id(state.gen_table_id_wrapped())
158            .to_internal_table_prost();
159
160        Ok(NodeBody::GapFill(Box::new(GapFillNode {
161            pointer_key_indices: self
162                .pointer_key_indices()
163                .into_iter()
164                .map(|idx| idx as u32)
165                .collect(),
166            time_column_index: self.time_col().index() as u32,
167            interval: Some(self.interval().to_expr_proto_checked_pure(
168                self.stream_kind().is_retract(),
169                "gap filling interval",
170            )?),
171            fill_columns: self
172                .fill_strategies()
173                .iter()
174                .map(|strategy| strategy.target_col.index() as u32)
175                .collect(),
176            fill_strategies,
177            state_table: Some(state_table),
178            partition_by_indices: self
179                .core
180                .partition_by_cols
181                .iter()
182                .map(|c| c.index() as u32)
183                .collect(),
184        })))
185    }
186}
187
188impl ExprRewritable<Stream> for StreamGapFill {
189    fn has_rewritable_expr(&self) -> bool {
190        true
191    }
192
193    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef<Stream> {
194        let mut core = self.core.clone();
195        core.rewrite_exprs(r);
196        Self {
197            base: self.base.clone_with_new_plan_id(),
198            core,
199        }
200        .into()
201    }
202}
203
204impl ExprVisitable for StreamGapFill {
205    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
206        self.core.visit_exprs(v)
207    }
208}
risingwave_frontend/optimizer/plan_node/stream_gap_fill.rs

risingwave_frontend/optimizer/plan_node/
stream_gap_fill.rs
