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

risingwave_frontend/optimizer/plan_node/
stream_iceberg_with_pk_index_writer.rs
