risingwave_frontend/optimizer/plan_node/
stream_iceberg_with_pk_index_writer.rs

1// Copyright 2026 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 anyhow::Context;
16use pretty_xmlish::{Pretty, XmlNode};
17use risingwave_common::catalog::{Field, Schema};
18use risingwave_common::types::DataType;
19use risingwave_common::util::sort_util::OrderType;
20use risingwave_connector::sink::catalog::desc::SinkDesc;
21use risingwave_pb::stream_plan::compaction_resolver_node::PkColumn;
22use risingwave_pb::stream_plan::stream_node::{NodeBody, PbStreamKind};
23use risingwave_pb::stream_plan::{
24    CompactionResolverNode, DispatchStrategy, DispatcherType, ExchangeNode,
25    IcebergWithPkIndexWriterNode, PbDispatchOutputMapping, PbStreamNode,
26};
27
28use super::stream::prelude::*;
29use crate::TableCatalog;
30use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
31use crate::optimizer::plan_node::utils::{Distill, TableCatalogBuilder, childless_record};
32use crate::optimizer::plan_node::{
33    ExprRewritable, PlanBase, PlanTreeNodeUnary, Stream, StreamExchange, StreamNode,
34    StreamPlanRef as PlanRef,
35};
36use crate::optimizer::property::{
37    Distribution, FunctionalDependencySet, MonotonicityMap, WatermarkColumns,
38};
39use crate::scheduler::SchedulerResult;
40use crate::stream_fragmenter::BuildFragmentGraphState;
41
42/// `StreamIcebergWithPkIndexWriter` is the stateful writer executor for the Iceberg
43/// with pk index sink. It maintains a PK index and writes data files to Iceberg.
44///
45/// Output schema: `[file_path: Varchar, position: Int64]` — delete-position info for
46/// the downstream `PositionDeleteMerger`.
47#[derive(Debug, Clone, PartialEq, Eq, Hash)]
48pub struct StreamIcebergWithPkIndexWriter {
49    pub base: PlanBase<Stream>,
50    pub input: PlanRef,
51    pub sink_desc: SinkDesc,
52    pub pk_index_table: TableCatalog,
53}
54
55impl StreamIcebergWithPkIndexWriter {
56    pub fn from_stream_sink(sink: &super::StreamSink) -> Result<Self> {
57        let output_schema = output_schema();
58        let fd_set = FunctionalDependencySet::new(output_schema.len());
59        let dist = match sink.distribution() {
60            Distribution::Single => Distribution::Single,
61            _ => Distribution::SomeShard,
62        };
63        let base = PlanBase::new_stream(
64            sink.ctx(),
65            output_schema,
66            sink.stream_key().map(|v| v.to_vec()),
67            fd_set,
68            dist,
69            StreamKind::AppendOnly,
70            sink.emit_on_window_close(),
71            WatermarkColumns::new(),
72            MonotonicityMap::new(),
73        );
74        let pk_index_table = build_iceberg_pk_state_table(sink.sink_desc())?;
75        Ok(Self {
76            base,
77            input: sink.input(),
78            sink_desc: sink.sink_desc().clone(),
79            pk_index_table,
80        })
81    }
82}
83
84fn output_schema() -> Schema {
85    Schema::new(vec![
86        Field::with_name(DataType::Varchar, "file_path"),
87        Field::with_name(DataType::Int64, "position"),
88    ])
89}
90
91/// Schema of the transient compaction resolver output consumed by the writer's right input.
92fn resolver_output_schema(sink_desc: &SinkDesc) -> Result<Schema> {
93    let downstream_pk = sink_desc
94        .downstream_pk
95        .as_deref()
96        .context("Missing downstream PK in Iceberg sink desc")?;
97    let mut fields: Vec<Field> = downstream_pk
98        .iter()
99        .map(|&idx| Field::from(&sink_desc.columns[idx].column_desc))
100        .collect();
101    fields.push(Field::with_name(DataType::Varchar, "file_path"));
102    fields.push(Field::with_name(DataType::Int64, "position"));
103    Ok(Schema::new(fields))
104}
105
106fn build_iceberg_pk_state_table(sink_desc: &SinkDesc) -> Result<TableCatalog> {
107    let mut builder = TableCatalogBuilder::default();
108
109    let downstream_pk = sink_desc
110        .downstream_pk
111        .as_deref()
112        .context("Missing downstream PK in Iceberg sink desc")?;
113    for &idx in downstream_pk {
114        builder.add_column(&Field::from(&sink_desc.columns[idx].column_desc));
115    }
116    builder.add_column(&Field::with_name(DataType::Varchar, "file_path"));
117    builder.add_column(&Field::with_name(DataType::Int64, "position"));
118
119    for idx in 0..downstream_pk.len() {
120        builder.add_order_column(idx, OrderType::ascending());
121    }
122
123    let res = builder.build((0..downstream_pk.len()).collect(), downstream_pk.len());
124    Ok(res)
125}
126
127impl Distill for StreamIcebergWithPkIndexWriter {
128    fn distill<'a>(&self) -> XmlNode<'a> {
129        let column_names = self
130            .sink_desc
131            .columns
132            .iter()
133            .map(|col| col.name_with_hidden().to_string())
134            .map(Pretty::from)
135            .collect();
136        let column_names = Pretty::Array(column_names);
137        let mut vec = Vec::with_capacity(2);
138        vec.push(("columns", column_names));
139        if let Some(pk) = &self.sink_desc.downstream_pk {
140            let column_names = pk
141                .iter()
142                .map(|&idx| self.sink_desc.columns[idx].name_with_hidden().to_string())
143                .map(Pretty::from)
144                .collect();
145            let column_names = Pretty::Array(column_names);
146            vec.push(("downstream_pk", column_names));
147        }
148
149        childless_record("StreamIcebergWithPkIndexWriter", vec)
150    }
151}
152
153impl PlanTreeNodeUnary<Stream> for StreamIcebergWithPkIndexWriter {
154    fn input(&self) -> PlanRef {
155        self.input.clone()
156    }
157
158    fn clone_with_input(&self, input: PlanRef) -> Self {
159        Self {
160            base: self.base.clone(),
161            input,
162            sink_desc: self.sink_desc.clone(),
163            pk_index_table: self.pk_index_table.clone(),
164        }
165    }
166}
167
168impl_plan_tree_node_for_unary! { Stream, StreamIcebergWithPkIndexWriter }
169
170impl StreamNode for StreamIcebergWithPkIndexWriter {
171    fn to_stream_prost_body(&self, _state: &mut BuildFragmentGraphState) -> NodeBody {
172        unreachable!(
173            "iceberg pk-index writer cannot be converted into a prost body -- call \
174             `adhoc_to_stream_prost` instead, since it declares a compaction resolver input."
175        )
176    }
177}
178
179impl StreamIcebergWithPkIndexWriter {
180    /// Serializes the writer with its normal upstream and a task-independent compaction resolver.
181    /// The exchange makes the resolver a separate fragment whose actors can remain inactive until
182    /// meta attaches them to the writer for a compaction task.
183    pub fn adhoc_to_stream_prost(
184        &self,
185        state: &mut BuildFragmentGraphState,
186    ) -> SchedulerResult<PbStreamNode> {
187        let pk_index_table = self
188            .pk_index_table
189            .clone()
190            .with_id(state.gen_table_id_wrapped())
191            .to_internal_table_prost();
192        let resolver_fields = resolver_output_schema(&self.sink_desc)
193            .context("build compaction resolver output schema")?
194            .to_prost();
195        let downstream_pk = self
196            .sink_desc
197            .downstream_pk
198            .as_ref()
199            .expect("validated when building the resolver schema");
200        let pk_columns: Vec<PkColumn> = downstream_pk
201            .iter()
202            .map(|&idx| PkColumn {
203                data_file_index: idx as u32,
204                column_desc: Some(self.sink_desc.columns[idx].column_desc.to_protobuf()),
205            })
206            .collect();
207        // The resolver output projects the PK columns first in `pk_columns` order, followed by
208        // `file_path` and `position`. These are output-schema indices, not data-file indices.
209        let resolver_stream_key: Vec<u32> = (0..pk_columns.len() as u32).collect();
210
211        // The writer's normal input must have a fragment boundary so that it is rendered as a
212        // merge and can be disconnected while applying compaction.
213        let left_input = if self.input.as_stream_exchange().is_some() {
214            self.input.to_stream_prost(state)?
215        } else {
216            let exchange: PlanRef = StreamExchange::new_no_shuffle(self.input.clone()).into();
217            exchange.to_stream_prost(state)?
218        };
219        let right_dispatcher = match self.distribution() {
220            Distribution::Single => DispatcherType::Simple,
221            _ => DispatcherType::Hash,
222        };
223        let resolver = PbStreamNode {
224            node_body: Some(NodeBody::CompactionResolver(Box::new(
225                CompactionResolverNode {
226                    sink_id: self.sink_desc.id,
227                    properties: self.sink_desc.properties.clone(),
228                    secret_refs: self.sink_desc.secret_refs.clone(),
229                    pk_columns,
230                },
231            ))),
232            identity: "CompactionResolver".into(),
233            operator_id: state.gen_operator_id(),
234            stream_key: resolver_stream_key.clone(),
235            fields: resolver_fields.clone(),
236            stream_kind: PbStreamKind::AppendOnly as i32,
237            ..Default::default()
238        };
239        let right_input = PbStreamNode {
240            node_body: Some(NodeBody::Exchange(Box::new(ExchangeNode {
241                strategy: Some(DispatchStrategy {
242                    r#type: right_dispatcher as i32,
243                    dist_key_indices: match right_dispatcher {
244                        DispatcherType::Hash => resolver_stream_key.clone(),
245                        DispatcherType::Simple => vec![],
246                        DispatcherType::Unspecified
247                        | DispatcherType::Broadcast
248                        | DispatcherType::NoShuffle => unreachable!(),
249                    },
250                    output_mapping: Some(PbDispatchOutputMapping::identical(resolver_fields.len())),
251                }),
252            }))),
253            input: vec![resolver],
254            identity: "IcebergCompactionResolverExchange".into(),
255            operator_id: state.gen_operator_id(),
256            stream_key: resolver_stream_key,
257            fields: resolver_fields,
258            stream_kind: PbStreamKind::AppendOnly as i32,
259        };
260
261        Ok(PbStreamNode {
262            node_body: Some(NodeBody::IcebergWithPkIndexWriter(Box::new(
263                IcebergWithPkIndexWriterNode {
264                    sink_desc: Some(self.sink_desc.to_proto()),
265                    pk_index_table: Some(pk_index_table),
266                },
267            ))),
268            input: vec![left_input, right_input],
269            identity: self.distill_to_string(),
270            operator_id: self.id().to_stream_node_operator_id(),
271            stream_key: self
272                .stream_key()
273                .unwrap_or_default()
274                .iter()
275                .map(|x| *x as u32)
276                .collect(),
277            fields: self.schema().to_prost(),
278            stream_kind: self.stream_kind().to_protobuf() as i32,
279        })
280    }
281}
282
283impl ExprRewritable<Stream> for StreamIcebergWithPkIndexWriter {}
284
285impl ExprVisitable for StreamIcebergWithPkIndexWriter {}
286
287#[cfg(test)]
288mod tests {
289    use std::collections::BTreeMap;
290
291    use risingwave_common::catalog::{
292        ColumnCatalog, ColumnDesc, ColumnId, CreateType, DEFAULT_SUPER_USER_ID, StreamJobStatus,
293    };
294    use risingwave_common::util::sort_util::ColumnOrder;
295    use risingwave_connector::sink::catalog::{SinkId, SinkType};
296
297    use super::*;
298
299    fn test_sink_desc() -> SinkDesc {
300        SinkDesc {
301            id: SinkId::placeholder(),
302            name: "s".to_owned(),
303            definition: "".to_owned(),
304            columns: vec![
305                ColumnCatalog::visible(ColumnDesc::named("id", ColumnId::new(0), DataType::Int32)),
306                ColumnCatalog::visible(ColumnDesc::named("v1", ColumnId::new(1), DataType::Int32)),
307                ColumnCatalog::hidden(ColumnDesc::named(
308                    "_row_id",
309                    ColumnId::new(2),
310                    DataType::Serial,
311                )),
312            ],
313            plan_pk: vec![ColumnOrder::new(2, OrderType::ascending())],
314            downstream_pk: Some(vec![1]),
315            distribution_key: vec![1],
316            properties: BTreeMap::new(),
317            secret_refs: BTreeMap::new(),
318            sink_type: SinkType::Upsert,
319            ignore_delete: false,
320            format_desc: None,
321            db_name: "dev".to_owned(),
322            sink_from_name: "t".to_owned(),
323            target_table: None,
324            extra_partition_col_idx: None,
325            create_type: CreateType::Foreground,
326            is_exactly_once: None,
327            auto_refresh_schema_from_table: None,
328        }
329    }
330
331    #[test]
332    fn test_build_iceberg_pk_state_table_uses_downstream_pk_columns() {
333        let table = build_iceberg_pk_state_table(&test_sink_desc()).unwrap();
334
335        assert_eq!(table.columns()[0].name(), "v1");
336        assert_eq!(table.columns()[1].name(), "file_path");
337        assert_eq!(table.columns()[2].name(), "position");
338        assert_eq!(table.pk().len(), 1);
339        assert_eq!(table.pk()[0].column_index, 0);
340        assert_eq!(table.distribution_key(), &[0]);
341        assert_eq!(table.read_prefix_len_hint, 1);
342        assert_eq!(table.owner, DEFAULT_SUPER_USER_ID);
343        assert_eq!(table.stream_job_status, StreamJobStatus::Creating);
344    }
345
346    #[test]
347    fn test_build_iceberg_pk_state_table_with_multi_column_pk() {
348        // Simulate planner output where the derived pk spans several stream-key columns, e.g. a
349        // hidden upstream column (`order_id`) promoted to visible and carried in verbatim.
350        let mut desc = test_sink_desc();
351        desc.columns.push(ColumnCatalog::visible(ColumnDesc::named(
352            "order_id",
353            ColumnId::new(3),
354            DataType::Int64,
355        )));
356        desc.columns.push(ColumnCatalog::visible(ColumnDesc::named(
357            "shard_id",
358            ColumnId::new(4),
359            DataType::Int64,
360        )));
361        desc.downstream_pk = Some(vec![1, 3, 4]); // v1, order_id, shard_id
362
363        let table = build_iceberg_pk_state_table(&desc).unwrap();
364
365        let names: Vec<_> = table
366            .columns()
367            .iter()
368            .map(|c| c.name().to_owned())
369            .collect();
370        assert_eq!(
371            names,
372            vec!["v1", "order_id", "shard_id", "file_path", "position"]
373        );
374        assert_eq!(table.pk().len(), 3);
375        assert_eq!(table.distribution_key(), &[0, 1, 2]);
376        assert_eq!(table.read_prefix_len_hint, 3);
377    }
378
379    #[test]
380    fn test_build_iceberg_pk_state_table_with_visible_extra_only() {
381        // Simulate planner output: downstream_pk = [user_pk, visible_extra].
382        let mut desc = test_sink_desc();
383        desc.columns.push(ColumnCatalog::visible(ColumnDesc::named(
384            "order_id",
385            ColumnId::new(3),
386            DataType::Int64,
387        )));
388        desc.downstream_pk = Some(vec![1, 3]); // v1, order_id
389
390        let table = build_iceberg_pk_state_table(&desc).unwrap();
391
392        let names: Vec<_> = table
393            .columns()
394            .iter()
395            .map(|c| c.name().to_owned())
396            .collect();
397        assert_eq!(names, vec!["v1", "order_id", "file_path", "position"]);
398        assert_eq!(table.pk().len(), 2);
399        assert_eq!(table.distribution_key(), &[0, 1]);
400        assert_eq!(table.read_prefix_len_hint, 2);
401    }
402}
risingwave_frontend/optimizer/plan_node/stream_iceberg_with_pk_index_writer.rs

risingwave_frontend/optimizer/plan_node/
stream_iceberg_with_pk_index_writer.rs
