risingwave_stream/executor/source/
iceberg_list_executor.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 std::sync::Arc;
16
17use anyhow::anyhow;
18use either::Either;
19use futures_async_stream::try_stream;
20use parking_lot::Mutex;
21use risingwave_common::array::Op;
22use risingwave_common::catalog::ColumnCatalog;
23use risingwave_common::config::StreamingConfig;
24use risingwave_common::system_param::local_manager::SystemParamsReaderRef;
25use risingwave_connector::source::ConnectorProperties;
26use risingwave_connector::source::iceberg::{
27    IcebergIncrementalScan, IcebergScanMetricsLabels, IcebergScanPlanner, IcebergScanProjection,
28    PersistedFileScanTask,
29};
30use risingwave_connector::source::reader::desc::SourceDescBuilder;
31use thiserror_ext::AsReport;
32use tokio::sync::mpsc::UnboundedReceiver;
33
34use super::{StreamSourceCore, barrier_to_message_stream};
35use crate::executor::prelude::*;
36use crate::executor::stream_reader::StreamReaderWithPause;
37
38pub struct IcebergListExecutor<S: StateStore> {
39    actor_ctx: ActorContextRef,
40
41    /// Streaming source for external
42    stream_source_core: StreamSourceCore<S>,
43
44    /// Columns of fetch executor, used to plan files.
45    /// For backward compatibility, this can be None, meaning all columns are needed.
46    downstream_columns: Option<Vec<ColumnCatalog>>,
47
48    /// Metrics for monitor.
49    #[expect(dead_code)]
50    metrics: Arc<StreamingMetrics>,
51
52    /// Receiver of barrier channel.
53    barrier_receiver: Option<UnboundedReceiver<Barrier>>,
54
55    /// System parameter reader to read barrier interval
56    #[expect(dead_code)]
57    system_params: SystemParamsReaderRef,
58
59    /// Rate limit in rows/s.
60    #[expect(dead_code)]
61    rate_limit_rps: Option<u32>,
62
63    /// Streaming config
64    streaming_config: Arc<StreamingConfig>,
65}
66
67impl<S: StateStore> IcebergListExecutor<S> {
68    #[expect(clippy::too_many_arguments)]
69    pub fn new(
70        actor_ctx: ActorContextRef,
71        stream_source_core: StreamSourceCore<S>,
72        downstream_columns: Option<Vec<ColumnCatalog>>,
73        metrics: Arc<StreamingMetrics>,
74        barrier_receiver: UnboundedReceiver<Barrier>,
75        system_params: SystemParamsReaderRef,
76        rate_limit_rps: Option<u32>,
77        streaming_config: Arc<StreamingConfig>,
78    ) -> Self {
79        Self {
80            actor_ctx,
81            stream_source_core,
82            downstream_columns,
83            metrics,
84            barrier_receiver: Some(barrier_receiver),
85            system_params,
86            rate_limit_rps,
87            streaming_config,
88        }
89    }
90
91    #[try_stream(ok = Message, error = StreamExecutorError)]
92    async fn into_stream(mut self) {
93        let mut barrier_receiver = self.barrier_receiver.take().unwrap();
94        let first_barrier = barrier_receiver
95            .recv()
96            .instrument_await("source_recv_first_barrier")
97            .await
98            .ok_or_else(|| {
99                anyhow!(
100                    "failed to receive the first barrier, actor_id: {:?}, source_id: {:?}",
101                    self.actor_ctx.id,
102                    self.stream_source_core.source_id
103                )
104            })?;
105        let first_epoch = first_barrier.epoch;
106
107        // Build source description from the builder.
108        let source_desc_builder: SourceDescBuilder =
109            self.stream_source_core.source_desc_builder.take().unwrap();
110
111        let properties = source_desc_builder.with_properties();
112        let config = ConnectorProperties::extract(properties, false)?;
113        let ConnectorProperties::Iceberg(iceberg_properties) = config else {
114            unreachable!()
115        };
116
117        let scan_projection =
118            IcebergScanProjection::from_downstream_columns(self.downstream_columns.as_deref());
119
120        tracing::debug!("scan_projection: {:?}", scan_projection);
121
122        yield Message::Barrier(first_barrier);
123        let barrier_stream = barrier_to_message_stream(barrier_receiver).boxed();
124
125        let source_id_str = self.stream_source_core.source_id.to_string();
126        let source_name_str = self.stream_source_core.source_name.clone();
127        let table_name = iceberg_properties.table.table_name().to_owned();
128        let scan_metrics = IcebergScanMetricsLabels::new(
129            source_id_str.clone(),
130            source_name_str.clone(),
131            table_name,
132        );
133        let scan_planner = IcebergScanPlanner::new(
134            (*iceberg_properties).clone(),
135            scan_projection,
136            Some(scan_metrics.clone()),
137        );
138
139        let state_table = self.stream_source_core.split_state_store.state_table_mut();
140        state_table.init_epoch(first_epoch).await?;
141        let state_row = state_table.get_from_one_value_table().await?;
142        // last_snapshot is EXCLUSIVE (i.e., already scanned)
143        let mut last_snapshot: Option<i64> = state_row.map(|s| *s.as_int64());
144        let mut prev_persisted_snapshot = last_snapshot;
145
146        if last_snapshot.is_none() {
147            // do a regular scan, then switch to incremental scan
148            // TODO: we may support starting from a specific snapshot/timestamp later
149            // If the current snapshot is None (empty table), go to incremental scan directly.
150            if let Some(snapshot_plan) = scan_planner.plan_current_snapshot().await? {
151                last_snapshot = Some(snapshot_plan.snapshot_id);
152                let mut chunk_builder = StreamChunkBuilder::new(
153                    self.streaming_config.developer.chunk_size,
154                    vec![DataType::Varchar, DataType::Jsonb],
155                );
156                #[for_await]
157                for scan_task in snapshot_plan.tasks {
158                    let scan_task = scan_task?;
159                    let data_file_path = scan_task.data_file_path.clone();
160                    let persisted_task = PersistedFileScanTask::encode(scan_task)?;
161                    if let Some(chunk) = chunk_builder.append_row(
162                        Op::Insert,
163                        &[
164                            Some(ScalarImpl::Utf8(data_file_path.into())),
165                            Some(ScalarImpl::Jsonb(persisted_task)),
166                        ],
167                    ) {
168                        yield Message::Chunk(chunk);
169                    }
170                }
171                if let Some(chunk) = chunk_builder.take() {
172                    yield Message::Chunk(chunk);
173                }
174            }
175        }
176
177        let last_snapshot = Arc::new(Mutex::new(last_snapshot));
178        let build_incremental_stream = || {
179            incremental_scan_stream(
180                scan_planner.clone(),
181                last_snapshot.clone(),
182                self.streaming_config.developer.iceberg_list_interval_sec,
183            )
184            .map(|res| match res {
185                Ok(scan_task) => {
186                    let data_file_path = scan_task.data_file_path.clone();
187                    let persisted_task = PersistedFileScanTask::encode(scan_task)?;
188                    let row = (
189                        Op::Insert,
190                        OwnedRow::new(vec![
191                            Some(ScalarImpl::Utf8(data_file_path.into())),
192                            Some(ScalarImpl::Jsonb(persisted_task)),
193                        ]),
194                    );
195                    Ok(StreamChunk::from_rows(
196                        &[row],
197                        &[DataType::Varchar, DataType::Jsonb],
198                    ))
199                }
200                Err(e) => Err(e),
201            })
202        };
203
204        let mut stream =
205            StreamReaderWithPause::<true, _>::new(barrier_stream, build_incremental_stream());
206
207        // TODO: support pause (incl. pause on startup)/resume/rate limit
208
209        while let Some(msg) = stream.next().await {
210            match msg {
211                Err(e) => {
212                    tracing::warn!(
213                        error = %e.as_report(),
214                        "incremental iceberg list stream errored, rebuilding"
215                    );
216                    scan_metrics.record_scan_error("list_error");
217                    stream.replace_data_stream(build_incremental_stream());
218                }
219                Ok(msg) => match msg {
220                    // Barrier arrives.
221                    Either::Left(msg) => match &msg {
222                        Message::Barrier(barrier) => {
223                            if let Some(mutation) = barrier.mutation.as_deref() {
224                                match mutation {
225                                    Mutation::Pause => stream.pause_stream(),
226                                    Mutation::Resume => stream.resume_stream(),
227                                    _ => (),
228                                }
229                            }
230                            if let Some(last_snapshot) = *last_snapshot.lock() {
231                                let state_row =
232                                    OwnedRow::new(vec![ScalarImpl::Int64(last_snapshot).into()]);
233                                if let Some(prev_persisted_snapshot_id) = prev_persisted_snapshot {
234                                    let prev_state_row = OwnedRow::new(vec![
235                                        ScalarImpl::Int64(prev_persisted_snapshot_id).into(),
236                                    ]);
237                                    state_table.update(prev_state_row, state_row);
238                                } else {
239                                    state_table.insert(state_row);
240                                }
241                                prev_persisted_snapshot = Some(last_snapshot);
242                            }
243                            state_table
244                                .commit_assert_no_update_vnode_bitmap(barrier.epoch)
245                                .await?;
246                            // Propagate the barrier.
247                            yield msg;
248                        }
249                        // Only barrier can be received.
250                        _ => unreachable!(),
251                    },
252                    // Data arrives.
253                    Either::Right(chunk) => {
254                        yield Message::Chunk(chunk);
255                    }
256                },
257            }
258        }
259    }
260}
261
262/// `last_snapshot` is EXCLUSIVE (i.e., already scanned)
263#[try_stream(
264    boxed,
265    ok = iceberg::scan::FileScanTask,
266    error = StreamExecutorError
267)]
268async fn incremental_scan_stream(
269    scan_planner: IcebergScanPlanner,
270    last_snapshot_lock: Arc<Mutex<Option<i64>>>,
271    list_interval_sec: u64,
272) {
273    let mut last_snapshot: Option<i64> = *last_snapshot_lock.lock();
274    loop {
275        tokio::time::sleep(std::time::Duration::from_secs(list_interval_sec)).await;
276
277        match scan_planner.plan_incremental(last_snapshot).await? {
278            IcebergIncrementalScan::EmptyTable | IcebergIncrementalScan::UpToDate { .. } => {}
279            IcebergIncrementalScan::Planned(plan) => {
280                #[for_await]
281                for scan_task in plan.tasks {
282                    yield scan_task?;
283                }
284
285                last_snapshot = Some(plan.snapshot_id);
286                *last_snapshot_lock.lock() = last_snapshot;
287                scan_planner.record_caught_up();
288            }
289        }
290    }
291}
292
293impl<S: StateStore> Execute for IcebergListExecutor<S> {
294    fn execute(self: Box<Self>) -> BoxedMessageStream {
295        self.into_stream().boxed()
296    }
297}
298
299impl<S: StateStore> Debug for IcebergListExecutor<S> {
300    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
301        f.debug_struct("IcebergListExecutor")
302            .field("source_id", &self.stream_source_core.source_id)
303            .field("column_ids", &self.stream_source_core.column_ids)
304            .finish()
305    }
306}
risingwave_stream/executor/source/iceberg_list_executor.rs

risingwave_stream/executor/source/
iceberg_list_executor.rs
