Skip to main content

risingwave_frontend/optimizer/plan_node/
stream_eowc_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::TableCatalog;
25use crate::binder::BoundFillStrategy;
26use crate::expr::{Expr, ExprImpl, ExprRewriter, ExprVisitor, InputRef};
27use crate::optimizer::plan_node::stream::StreamPlanNodeMetadata;
28use crate::optimizer::plan_node::utils::impl_distill_by_unit;
29use crate::optimizer::property::{Distribution, MonotonicityMap, WatermarkColumns};
30use crate::scheduler::SchedulerResult;
31use crate::stream_fragmenter::BuildFragmentGraphState;
32
33/// `StreamEowcGapFill` implements [`super::Stream`] to represent a gap-filling operation on a time
34/// series in streaming mode.
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub struct StreamEowcGapFill {
37    pub base: PlanBase<super::Stream>,
38    core: generic::GapFill<PlanRef<Stream>>,
39}
40
41impl StreamEowcGapFill {
42    pub fn new(core: generic::GapFill<PlanRef<Stream>>) -> Self {
43        let input = &core.input;
44
45        let dist = if core.partition_by_cols.is_empty() {
46            Distribution::Single
47        } else {
48            let partition_indices: Vec<usize> =
49                core.partition_by_cols.iter().map(|c| c.index()).collect();
50            Distribution::HashShard(partition_indices)
51        };
52
53        // A late anchor makes gap fill back-fill a partition below the already-forwarded time
54        // watermark, so the time watermark is never preserved. Only `PARTITION BY` columns are
55        // copied unchanged into generated rows, so only their watermark can be propagated.
56        let mut watermark_columns = WatermarkColumns::new();
57        for partition_col in &core.partition_by_cols {
58            let idx = partition_col.index();
59            if let Some(group) = input.watermark_columns().get_group(idx) {
60                watermark_columns.insert(idx, group);
61            }
62        }
63
64        // Without partitioning the output stays time-sorted, so the time column is monotonic.
65        // With partitioning, cross-partition interleaving breaks that; only the copied
66        // `PARTITION BY` columns keep their monotonicity.
67        let mut columns_monotonicity = MonotonicityMap::new();
68        if core.partition_by_cols.is_empty() {
69            let idx = core.time_col.index();
70            columns_monotonicity.insert(idx, input.columns_monotonicity()[idx]);
71        } else {
72            for partition_col in &core.partition_by_cols {
73                let idx = partition_col.index();
74                columns_monotonicity.insert(idx, input.columns_monotonicity()[idx]);
75            }
76        }
77
78        let base = PlanBase::new_stream_with_core(
79            &core,
80            dist,
81            input.stream_kind(),
82            true, // provides EOWC semantics
83            watermark_columns,
84            columns_monotonicity,
85        );
86        Self { base, core }
87    }
88
89    pub fn time_col(&self) -> &InputRef {
90        &self.core.time_col
91    }
92
93    pub fn interval(&self) -> &ExprImpl {
94        &self.core.interval
95    }
96
97    pub fn fill_strategies(&self) -> &[BoundFillStrategy] {
98        &self.core.fill_strategies
99    }
100
101    fn infer_prev_row_table(&self) -> TableCatalog {
102        let mut tbl_builder = TableCatalogBuilder::default();
103
104        for field in self.core.schema().fields() {
105            tbl_builder.add_column(field);
106        }
107
108        if self.core.partition_by_cols.is_empty() {
109            // No partition: single-row table, use first column as PK
110            if !self.core.schema().fields().is_empty() {
111                tbl_builder.add_order_column(0, OrderType::ascending());
112            }
113            tbl_builder.build(vec![], 0)
114        } else {
115            // With partition: PK = partition_cols (one prev_row per partition)
116            for pc in &self.core.partition_by_cols {
117                tbl_builder.add_order_column(pc.index(), OrderType::ascending());
118            }
119            let dist_key_indices: Vec<usize> = self
120                .core
121                .partition_by_cols
122                .iter()
123                .map(|c| c.index())
124                .collect();
125            tbl_builder.build(dist_key_indices, 0)
126        }
127    }
128}
129
130impl PlanTreeNodeUnary<Stream> for StreamEowcGapFill {
131    fn input(&self) -> PlanRef<Stream> {
132        self.core.input.clone()
133    }
134
135    fn clone_with_input(&self, input: PlanRef<Stream>) -> Self {
136        let mut core = self.core.clone();
137        core.input = input;
138        Self::new(core)
139    }
140}
141
142impl_plan_tree_node_for_unary! { Stream, StreamEowcGapFill }
143impl_distill_by_unit!(StreamEowcGapFill, core, "StreamEowcGapFill");
144
145impl TryToStreamPb for StreamEowcGapFill {
146    fn try_to_stream_prost_body(
147        &self,
148        state: &mut BuildFragmentGraphState,
149    ) -> SchedulerResult<NodeBody> {
150        use risingwave_pb::stream_plan::*;
151
152        let fill_strategies: Vec<String> = self
153            .fill_strategies()
154            .iter()
155            .map(|strategy| match strategy.strategy {
156                crate::binder::FillStrategy::Locf => "locf".to_owned(),
157                crate::binder::FillStrategy::Interpolate => "interpolate".to_owned(),
158                crate::binder::FillStrategy::Null => "null".to_owned(),
159            })
160            .collect();
161
162        let prev_row_table = self
163            .infer_prev_row_table()
164            .with_id(state.gen_table_id_wrapped())
165            .to_internal_table_prost();
166
167        Ok(NodeBody::EowcGapFill(Box::new(EowcGapFillNode {
168            time_column_index: self.time_col().index() as u32,
169            interval: Some(self.interval().to_expr_proto_checked_pure(
170                self.stream_kind().is_retract(),
171                "gap filling interval",
172            )?),
173            fill_columns: self
174                .fill_strategies()
175                .iter()
176                .map(|strategy| strategy.target_col.index() as u32)
177                .collect(),
178            fill_strategies,
179            prev_row_table: Some(prev_row_table),
180            partition_by_indices: self
181                .core
182                .partition_by_cols
183                .iter()
184                .map(|c| c.index() as u32)
185                .collect(),
186        })))
187    }
188}
189
190impl ExprRewritable<Stream> for StreamEowcGapFill {
191    fn has_rewritable_expr(&self) -> bool {
192        true
193    }
194
195    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef<Stream> {
196        let mut core = self.core.clone();
197        core.rewrite_exprs(r);
198        Self {
199            base: self.base.clone_with_new_plan_id(),
200            core,
201        }
202        .into()
203    }
204}
205
206impl ExprVisitable for StreamEowcGapFill {
207    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
208        self.core.visit_exprs(v)
209    }
210}