risingwave_frontend/optimizer/plan_node/
stream_eowc_gap_fill.rs1use 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#[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 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 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, 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 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 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}