risingwave_stream/from_proto/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 std::sync::Arc;
16
17use anyhow::anyhow;
18use risingwave_common::secret::LocalSecretManager;
19use risingwave_common::types::DataType;
20use risingwave_connector::sink::iceberg::{
21    ICEBERG_SINK, IcebergConfig, LEGACY_PK_INDEX_FILE_PATH, PkIndexDataFileLocator,
22    create_and_validate_table_impl,
23};
24use risingwave_connector::sink::{SinkMetaClient, SinkWriterParam};
25use risingwave_expr::bail;
26use risingwave_pb::id::SinkId;
27use risingwave_pb::stream_plan::IcebergWithPkIndexWriterNode;
28use risingwave_storage::StateStore;
29
30use super::super::sink::build_sink_param;
31use crate::common::table::state_table::{StateTableBuilder, StateTableOpConsistencyLevel};
32use crate::error::StreamResult;
33use crate::executor::{Executor, IcebergWriterImpl, StreamExecutorError, WriterExecutor};
34use crate::from_proto::ExecutorBuilder;
35use crate::task::ExecutorParams;
36pub struct IcebergWithPkIndexWriterExecutorBuilder;
37
38impl_stream_node_body!(IcebergWithPkIndexWriter(IcebergWithPkIndexWriterNode) => IcebergWithPkIndexWriterExecutorBuilder);
39
40impl ExecutorBuilder for IcebergWithPkIndexWriterExecutorBuilder {
41    type Node = IcebergWithPkIndexWriterNode;
42
43    async fn new_boxed_executor(
44        params: ExecutorParams,
45        node: &Self::Node,
46        store: impl StateStore,
47    ) -> StreamResult<Executor> {
48        let [input, resolver_input] = params
49            .input
50            .try_into()
51            .map_err(|_| anyhow!("IcebergWithPkIndexWriterExecutor requires exactly two inputs"))?;
52        let sink_desc = node.sink_desc.as_ref().unwrap();
53        let sink_id: SinkId = sink_desc.get_id();
54        let sink_name = sink_desc.get_name().to_owned();
55
56        let properties_with_secret = LocalSecretManager::global().fill_secrets(
57            sink_desc.get_properties().clone(),
58            sink_desc.get_secret_refs().clone(),
59        )?;
60        let config = IcebergConfig::from_btreemap(properties_with_secret.clone())
61            .map_err(|err| StreamExecutorError::from((err, sink_id)))?;
62
63        let pk_indices = sink_desc
64            .downstream_pk
65            .iter()
66            .map(|&idx| idx as usize)
67            .collect::<Vec<_>>();
68        if pk_indices.is_empty() {
69            bail!("missing downstream pk in iceberg sink desc");
70        }
71
72        let (sink_param, _) = build_sink_param(sink_desc, properties_with_secret, ICEBERG_SINK)?;
73
74        let table = create_and_validate_table_impl(&config, &sink_param)
75            .await
76            .map_err(|e| StreamExecutorError::sink_error(e, sink_id))?;
77
78        let pk_index_state_table = StateTableBuilder::new(
79            node.get_pk_index_table()?,
80            store,
81            Some(Arc::new(params.vnode_bitmap.clone())),
82        )
83        .enable_preload_all_rows_by_config(&params.config)
84        .with_op_consistency_level(StateTableOpConsistencyLevel::Inconsistent)
85        .build()
86        .await;
87
88        let meta_client = params
89            .env
90            .meta_client()
91            .ok_or_else(|| anyhow!("meta client is required for Iceberg writer"))?;
92        let meta_client = SinkMetaClient::MetaClient(meta_client);
93
94        let writer_param = SinkWriterParam {
95            executor_id: params.executor_id,
96            vnode_bitmap: params.vnode_bitmap.clone(),
97            meta_client: Some(meta_client),
98            extra_partition_col_idx: sink_desc.extra_partition_col_idx.map(|v| v as usize),
99            actor_id: params.actor_context.id,
100            sink_id,
101            sink_name,
102            connector: ICEBERG_SINK.to_owned(),
103            streaming_config: params.config.as_ref().clone(),
104            time_zone: params.actor_context.time_zone,
105        };
106        let data_file_locator = if config.pk_index_stores_file_uuid() {
107            Some(
108                PkIndexDataFileLocator::new(table.metadata())
109                    .map_err(|e| StreamExecutorError::sink_error(e, sink_id))?,
110            )
111        } else {
112            None
113        };
114        let data_types = pk_index_state_table.get_data_types();
115        let expected_data_file_type = match data_file_locator {
116            Some(_) => DataType::Bytea,
117            None => DataType::Varchar,
118        };
119        if data_types.len() < 2 || data_types[data_types.len() - 2] != expected_data_file_type {
120            bail!(
121                "iceberg pk-index state table {:?} does not match `{}` of sink {}",
122                data_types,
123                LEGACY_PK_INDEX_FILE_PATH,
124                sink_id
125            );
126        }
127        let writer = IcebergWriterImpl::build(&config, table, &writer_param)?;
128
129        let exec = WriterExecutor::new(
130            params.actor_context,
131            input,
132            resolver_input,
133            pk_indices,
134            pk_index_state_table,
135            writer,
136            data_file_locator,
137            params.config.developer.chunk_size,
138            params
139                .config
140                .developer
141                .iceberg_pk_index_sink_lookup_concurrency,
142            params.config.developer.iceberg_pk_index_sink_max_batch_rows,
143            sink_id,
144            params.local_barrier_manager.clone(),
145        );
146        Ok((params.info, exec).into())
147    }
148}
risingwave_stream/from_proto/iceberg_with_pk_index/writer.rs

risingwave_stream/from_proto/iceberg_with_pk_index/
writer.rs
