Skip to main content

risingwave_connector/source/iceberg/
update_reader.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//! Per-file updates (Insert and Delete) for PK-index sink output. The planner, not the reader's
16//! local history, determines whether a file is new or retained. This does not schedule phases or
17//! commit progress: callers must checkpoint cursors with output and enforce Delete-before-Insert.
18
19use anyhow::Context;
20use futures::StreamExt;
21use futures_async_stream::try_stream;
22use iceberg::delete_vector::DeleteVector;
23use iceberg::metadata_columns::RESERVED_FIELD_ID_POS;
24use iceberg::scan::{FileScanTask, FileScanTaskDeleteFile};
25use iceberg::spec::DataFileFormat;
26use iceberg::table::Table;
27use risingwave_common::array::arrow::IcebergArrowConvert;
28use risingwave_common::array::arrow::arrow_array_iceberg::Int64Array;
29use risingwave_common::array::{Op, StreamChunk};
30use risingwave_common::bitmap::{Bitmap, BitmapBuilder};
31use risingwave_common::ensure;
32
33use super::validate_update_projection;
34use crate::error::{ConnectorError, ConnectorResult};
35use crate::sink::iceberg::PositionDeleteReader;
36
37/// The planner must retain this classification and the parent descriptors across replay.
38#[derive(Clone)]
39pub enum IcebergUpdateReadMode {
40    /// Bootstrap or newly added file: current deletes only filter Inserts.
41    Insert,
42    /// A file present in both snapshots: retract only the new deleted positions.
43    Delete {
44        parent_deletes: Vec<FileScanTaskDeleteFile>,
45    },
46    /// Removed data file: retract its previously visible rows.
47    Remove,
48}
49
50/// An empty chunk still carries read progress. EOF, not the last nonempty chunk, completes a task.
51pub struct IcebergUpdateBatch {
52    pub chunk: StreamChunk,
53    pub next_position: u64,
54}
55
56pub fn newly_deleted_positions(
57    parent: &DeleteVector,
58    mut current: DeleteVector,
59) -> ConnectorResult<DeleteVector> {
60    ensure!(
61        parent.is_subset(&current),
62        "deleted positions shrank for a retained data file"
63    );
64    current -= parent;
65    Ok(current)
66}
67
68/// Read one whole immutable file with bounded row batches. SDK `_pos` is used internally and
69/// removed before creating logical chunks. No shredded-VARIANT-to-NULL fallback is allowed here.
70/// Resume rereads from the beginning; seeking is an optimization, not part of the cursor contract.
71/// The caller supplies a reader for this table's `FileIO` and controls its cache lifetime. Sequential
72/// tasks can reuse it after dropping the returned stream; parallel tasks need separate readers.
73#[try_stream(boxed, ok = IcebergUpdateBatch, error = ConnectorError)]
74pub async fn read_file_updates(
75    table: Table,
76    delete_reader: &mut PositionDeleteReader,
77    mut task: FileScanTask,
78    mode: IcebergUpdateReadMode,
79    chunk_size: usize,
80    resume_position: u64,
81) {
82    ensure!(chunk_size > 0, "update reader batch size must be positive");
83    ensure!(
84        task.data_file_format == DataFileFormat::Parquet,
85        "update reader requires Parquet data files"
86    );
87    ensure!(
88        task.key_metadata.is_none(),
89        "encrypted data files are not supported by the update reader"
90    );
91    ensure!(
92        task.start == 0 && task.length == task.file_size_in_bytes && task.predicate.is_none(),
93        "update reader requires a whole-file task without predicate pushdown"
94    );
95    let record_count = task
96        .record_count
97        .context("update task is missing data record_count")?;
98    ensure!(
99        record_count <= i64::MAX as u64 && resume_position <= record_count,
100        "invalid update task position"
101    );
102    validate_update_projection(&task.schema, &task.project_field_ids)?;
103
104    let current = delete_reader
105        .read_file_scoped(&task.data_file_path, record_count, &task.deletes)
106        .await?;
107    let newly_deleted = matches!(&mode, IcebergUpdateReadMode::Delete { .. });
108    let (op, positions) = match mode {
109        IcebergUpdateReadMode::Insert => (Op::Insert, current),
110        IcebergUpdateReadMode::Remove => (Op::Delete, current),
111        IcebergUpdateReadMode::Delete { parent_deletes } => {
112            let parent = delete_reader
113                .read_file_scoped(&task.data_file_path, record_count, &parent_deletes)
114                .await?;
115            (Op::Delete, newly_deleted_positions(&parent, current)?)
116        }
117    };
118    task.deletes.clear();
119    let logical_column_count = task.project_field_ids.len();
120    task.project_field_ids.push(RESERVED_FIELD_ID_POS);
121    let reader = table
122        .reader_builder()
123        .with_batch_size(chunk_size)
124        .with_data_file_concurrency_limit(1)
125        .with_row_group_filtering_enabled(false)
126        .build();
127    let mut batches = reader.read(tokio_stream::once(Ok(task)).boxed())?.stream();
128    let projection: Vec<_> = (0..logical_column_count).collect();
129    let mut next_position = 0u64;
130    while let Some(batch) = batches.next().await {
131        let batch = batch?;
132        ensure!(
133            batch.num_columns() == logical_column_count + 1,
134            "unexpected update reader output schema"
135        );
136        let row_positions = batch
137            .column(logical_column_count)
138            .as_any()
139            .downcast_ref::<Int64Array>()
140            .context("SDK row positions must be int64")?;
141        let mut visibility = BitmapBuilder::with_capacity(batch.num_rows());
142        for pos in row_positions {
143            let pos = u64::try_from(pos.context("null SDK row position")?)
144                .context("negative SDK row position")?;
145            ensure!(
146                pos == next_position && pos < record_count,
147                "SDK row positions disagree with the whole-file task"
148            );
149            next_position = pos + 1;
150            visibility.append(
151                pos >= resume_position
152                    && if newly_deleted {
153                        positions.contains(pos)
154                    } else {
155                        !positions.contains(pos)
156                    },
157            );
158        }
159        if next_position <= resume_position {
160            continue;
161        }
162        let mut visibility = visibility.finish();
163        let mut projected_batch = batch.project(&projection)?;
164        if visibility.count_ones() == 0 {
165            // Preserve the typed schema and physical cursor without converting filtered row values.
166            projected_batch = projected_batch.slice(0, 0);
167            visibility = Bitmap::zeros(0);
168        }
169        let mut chunk = IcebergArrowConvert.chunk_from_record_batch(&projected_batch)?;
170        chunk.set_visibility(visibility);
171        yield IcebergUpdateBatch {
172            chunk: StreamChunk::from_parts(vec![op; chunk.capacity()], chunk),
173            next_position,
174        };
175    }
176    ensure!(
177        next_position == record_count,
178        "data row count differs from the update task"
179    );
180}
181
182#[cfg(test)]
183#[path = "update_reader_test.rs"]
184mod tests;