risingwave_stream/executor/source/iceberg_update/
fetch.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
15//! Fetch reads assigned Iceberg update tasks and keeps a resumable cursor per task. Tasks are
16//! read in state-table order, so the tasks of one route are applied in dispatch order.
17
18use std::ops::Bound;
19use std::time::Instant;
20
21use anyhow::Context;
22use either::Either;
23use futures::stream;
24use risingwave_common::array::Op;
25use risingwave_common::catalog::ColumnCatalog;
26use risingwave_common::config::StreamingConfig;
27use risingwave_common::hash::VnodeBitmapExt;
28use risingwave_common::{bail, ensure};
29use risingwave_common_estimate_size::EstimateSize;
30use risingwave_connector::sink::iceberg::PositionDeleteReader;
31use risingwave_connector::source::iceberg::update_reader::read_file_updates;
32use risingwave_connector::source::iceberg::update_state::{
33    IcebergUpdateAssignment, IcebergUpdateFetchState,
34};
35use risingwave_connector::source::iceberg::{
36    GLOBAL_ICEBERG_SCAN_METRICS, IcebergFileScanMetrics, IcebergProperties,
37    IcebergScanMetricsLabels,
38};
39use risingwave_storage::store::PrefetchOptions;
40use thiserror_ext::AsReport;
41
42use super::{decode, properties, scan_metrics, task_pk, task_row};
43use crate::executor::prelude::*;
44use crate::executor::source::{SourceStateTableHandler, StreamSourceCore};
45use crate::executor::stream_reader::StreamReaderWithPause;
46
47pub struct IcebergUpdateFetchExecutor<S: StateStore> {
48    actor: ActorContextRef,
49    core: StreamSourceCore<S>,
50    upstream: Executor,
51    streaming_config: Arc<StreamingConfig>,
52    output_columns: Vec<ColumnCatalog>,
53    key_indices: Vec<usize>,
54}
55
56impl<S: StateStore> IcebergUpdateFetchExecutor<S> {
57    pub fn new(
58        actor: ActorContextRef,
59        core: StreamSourceCore<S>,
60        upstream: Executor,
61        streaming_config: Arc<StreamingConfig>,
62        output_columns: Vec<ColumnCatalog>,
63        key_indices: Vec<usize>,
64    ) -> Self {
65        Self {
66            actor,
67            core,
68            upstream,
69            streaming_config,
70            output_columns,
71            key_indices,
72        }
73    }
74
75    #[try_stream(ok = Message, error = StreamExecutorError)]
76    async fn into_stream(mut self) {
77        let developer = &self.streaming_config.developer;
78        let batch_size = developer.iceberg_fetch_batch_size as usize;
79        let chunk_size = developer.chunk_size;
80        ensure!(batch_size > 0, "Iceberg fetch batch size must be positive");
81        let properties = properties(&mut self.core)?;
82        let output_types = self
83            .output_columns
84            .iter()
85            .map(|column| column.data_type().clone())
86            .collect::<Vec<_>>();
87        let ctx = FetchContext {
88            job_id: self
89                .core
90                .split_state_store
91                .state_table()
92                .table_id()
93                .as_raw_id(),
94            batch_size,
95            chunk_size,
96            output_columns: self.output_columns,
97            key_indices: self.key_indices,
98            scan_metrics: scan_metrics(&self.core, &properties),
99            file_scan_metrics: IcebergFileScanMetrics::new(
100                &GLOBAL_ICEBERG_SCAN_METRICS,
101                properties.table.table_name(),
102            ),
103            properties,
104        };
105
106        let mut state_table = self.core.split_state_store;
107        let mut upstream = self.upstream.execute();
108        let first = expect_first_barrier(&mut upstream).await?;
109        let first_epoch = first.epoch;
110        let paused = first.is_pause_on_startup();
111        yield Message::Barrier(first);
112        state_table.init_epoch(first_epoch).await?;
113        let mut stream = FetchStream::new(upstream, stream::pending());
114        if paused {
115            stream.pause_stream();
116        }
117        let mut inflight = ctx.install_next_batch(&state_table, &mut stream).await?;
118        // Stays set until the installed batch reports `BatchFinished` or fails.
119        let mut busy = inflight > 0;
120        while let Some(message) = stream.next().await {
121            match message? {
122                Either::Left(Message::Barrier(barrier)) => {
123                    match barrier.mutation.as_deref() {
124                        Some(Mutation::Pause) => stream.pause_stream(),
125                        Some(Mutation::Resume) => stream.resume_stream(),
126                        _ => {}
127                    }
128                    // Rescale restarts the actor and restores its assigned vnodes at startup.
129                    barrier.assume_no_update_vnode_bitmap(self.actor.id)?;
130                    state_table.commit(barrier.epoch).await?;
131                    yield Message::Barrier(barrier);
132                    if !busy {
133                        inflight = ctx.install_next_batch(&state_table, &mut stream).await?;
134                        busy = inflight > 0;
135                    }
136                }
137                Either::Left(Message::Chunk(chunk)) => {
138                    ctx.accept_assignments(&mut state_table, &chunk).await?;
139                    if !busy {
140                        inflight = ctx.install_next_batch(&state_table, &mut stream).await?;
141                        busy = inflight > 0;
142                    }
143                }
144                Either::Left(_) => bail!("unexpected Iceberg Fetch watermark"),
145                Either::Right(Err(e)) => {
146                    tracing::error!(
147                        error = %e.as_report(),
148                        "Iceberg update Fetch errored, retrying"
149                    );
150                    ctx.scan_metrics.record_fetch_error();
151                    // Every emitted batch already saved its cursor, so the next barrier
152                    // resumes the unfinished tasks from local state.
153                    inflight = 0;
154                    ctx.scan_metrics.set_inflight_file_count(0);
155                    busy = false;
156                    stream.replace_data_stream(stream::pending());
157                }
158                Either::Right(Ok(FetchResult::Batch { state, chunk })) => {
159                    ensure!(
160                        chunk
161                            .columns()
162                            .iter()
163                            .map(|column| column.data_type())
164                            .collect::<Vec<_>>()
165                            == output_types,
166                        "Iceberg update output schema differs from its plan"
167                    );
168                    let table = state_table.state_table_mut();
169                    let old = table
170                        .get_row(task_pk(&state.assignment))
171                        .await?
172                        .context("missing Iceberg task state")?;
173                    table.update(old, task_row(&state.assignment, &state)?);
174                    if chunk.cardinality() != 0 {
175                        yield Message::Chunk(chunk);
176                    }
177                }
178                Either::Right(Ok(FetchResult::Finished(state))) => {
179                    // Retirement and the last output share the same checkpoint. Recovery
180                    // restores either the task cursor or its completed downstream effects.
181                    let table = state_table.state_table_mut();
182                    if let Some(old) = table.get_row(task_pk(&state.assignment)).await? {
183                        table.delete(old);
184                    }
185                    inflight = inflight.saturating_sub(1);
186                    ctx.scan_metrics.set_inflight_file_count(inflight);
187                }
188                Either::Right(Ok(FetchResult::BatchFinished)) => {
189                    inflight = ctx.install_next_batch(&state_table, &mut stream).await?;
190                    busy = inflight > 0;
191                }
192            }
193        }
194    }
195}
196
197impl<S: StateStore> Execute for IcebergUpdateFetchExecutor<S> {
198    fn execute(self: Box<Self>) -> BoxedMessageStream {
199        self.into_stream().boxed()
200    }
201}
202
203impl<S: StateStore> Debug for IcebergUpdateFetchExecutor<S> {
204    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
205        f.debug_struct("IcebergUpdateFetchExecutor")
206            .finish_non_exhaustive()
207    }
208}
209
210type FetchStream = StreamReaderWithPause<true, StreamExecutorResult<FetchResult>>;
211
212/// Inputs that stay fixed for the lifetime of a Fetch executor.
213struct FetchContext {
214    properties: IcebergProperties,
215    job_id: u32,
216    batch_size: usize,
217    chunk_size: usize,
218    output_columns: Vec<ColumnCatalog>,
219    key_indices: Vec<usize>,
220    scan_metrics: IcebergScanMetricsLabels,
221    file_scan_metrics: IcebergFileScanMetrics,
222}
223
224impl FetchContext {
225    /// Persists the assignments sent by List. A replayed assignment must match the stored one.
226    async fn accept_assignments<S: StateStore>(
227        &self,
228        state_table: &mut SourceStateTableHandler<S>,
229        chunk: &StreamChunk,
230    ) -> StreamExecutorResult<()> {
231        let table = state_table.state_table_mut();
232        for (op, row) in chunk.rows() {
233            let assignment: IcebergUpdateAssignment = decode(row, |assignment| assignment)?;
234            ensure!(
235                assignment.job_id == self.job_id,
236                "Iceberg assignment belongs to another job"
237            );
238            assignment
239                .task
240                .validate_key(&self.output_columns, &self.key_indices)?;
241            match (op, table.get_row(task_pk(&assignment)).await?) {
242                (Op::Insert, None) => {
243                    let state = IcebergUpdateFetchState::new(assignment);
244                    table.insert(task_row(&state.assignment, &state)?);
245                }
246                (Op::Insert, Some(old)) => {
247                    ensure!(
248                        decode(&old, |state: &IcebergUpdateFetchState| &state.assignment)?
249                            .assignment
250                            == assignment,
251                        "conflicting duplicate Iceberg assignment"
252                    );
253                }
254                _ => bail!("invalid Iceberg assignment operation"),
255            }
256        }
257        state_table.try_flush().await
258    }
259
260    /// Installs the next batch of unfinished tasks and returns its size. An empty batch
261    /// installs nothing, so an idle Fetch scans its state once per barrier instead of
262    /// again on `BatchFinished`.
263    async fn install_next_batch<S: StateStore>(
264        &self,
265        state_table: &SourceStateTableHandler<S>,
266        reader: &mut FetchStream,
267    ) -> StreamExecutorResult<usize> {
268        let tasks = pending_tasks(
269            state_table,
270            self.job_id,
271            self.batch_size,
272            &self.output_columns,
273            &self.key_indices,
274        )
275        .await?;
276        let count = tasks.len();
277        if tasks.is_empty() {
278            reader.replace_data_stream(stream::pending());
279        } else {
280            reader.replace_data_stream(
281                fetch_work(
282                    self.properties.clone(),
283                    tasks,
284                    self.chunk_size,
285                    self.file_scan_metrics.clone(),
286                )
287                .map(Ok),
288            );
289        }
290        self.scan_metrics.set_inflight_file_count(count);
291        Ok(count)
292    }
293}
294
295enum FetchResult {
296    Batch {
297        state: IcebergUpdateFetchState,
298        chunk: StreamChunk,
299    },
300    Finished(IcebergUpdateFetchState),
301    BatchFinished,
302}
303
304/// Reads a nonempty batch of tasks in order, then reports `BatchFinished`.
305#[try_stream(boxed, ok = FetchResult, error = StreamExecutorError)]
306async fn fetch_work(
307    properties: IcebergProperties,
308    tasks: Vec<IcebergUpdateFetchState>,
309    chunk_size: usize,
310    metrics: IcebergFileScanMetrics,
311) {
312    let table = properties.load_table().await?;
313    let mut reader = PositionDeleteReader::new(table.file_io());
314    for mut state in tasks {
315        let started = Instant::now();
316        let mut rows = 0;
317        let (scan, mode) = state.assignment.task.prepare_read(&table)?;
318        let mut batches = read_file_updates(
319            table.clone(),
320            &mut reader,
321            scan,
322            mode,
323            chunk_size,
324            state.next_position,
325        );
326        while let Some(batch) = batches.next().await {
327            let batch = batch?;
328            state.advance(batch.next_position, false)?;
329            rows += batch.chunk.cardinality() as u64;
330            metrics.record_read_bytes(batch.chunk.estimated_heap_size() as u64);
331            yield FetchResult::Batch {
332                state: state.clone(),
333                chunk: batch.chunk,
334            };
335        }
336        metrics.record_file_read_duration(started.elapsed().as_secs_f64());
337        if rows > 0 {
338            metrics.record_rows_read(rows);
339        }
340        metrics.record_file_read();
341        state.advance(state.assignment.task.record_count(), true)?;
342        yield FetchResult::Finished(state);
343    }
344    yield FetchResult::BatchFinished;
345}
346
347pub(super) async fn pending_tasks<S: StateStore>(
348    state: &SourceStateTableHandler<S>,
349    job_id: u32,
350    limit: usize,
351    columns: &[ColumnCatalog],
352    key_indices: &[usize],
353) -> StreamExecutorResult<Vec<IcebergUpdateFetchState>> {
354    ensure!(limit > 0, "Iceberg fetch batch size must be positive");
355    let table = state.state_table();
356    let mut tasks = vec![];
357    'vnodes: for vnode in table.vnodes().iter_vnodes() {
358        let iter = table
359            .iter_with_vnode(
360                vnode,
361                &(Bound::<OwnedRow>::Unbounded, Bound::<OwnedRow>::Unbounded),
362                PrefetchOptions::prefetch_for_small_range_scan(),
363            )
364            .await?;
365        pin_mut!(iter);
366        while let Some(row) = iter.next().await {
367            let task = decode(&row?, |task: &IcebergUpdateFetchState| &task.assignment)?;
368            ensure!(
369                task.assignment.job_id == job_id,
370                "Iceberg task belongs to another ingestion job"
371            );
372            task.assignment.task.validate_key(columns, key_indices)?;
373            ensure!(
374                !task.finished && task.next_position <= task.assignment.task.record_count(),
375                "invalid persisted Iceberg task cursor"
376            );
377            tasks.push(task);
378            if tasks.len() == limit {
379                break 'vnodes;
380            }
381        }
382    }
383    Ok(tasks)
384}
risingwave_stream/executor/source/iceberg_update/fetch.rs

risingwave_stream/executor/source/iceberg_update/
fetch.rs
