risingwave_connector/source/iceberg/
update_reader.rs1use 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#[derive(Clone)]
39pub enum IcebergUpdateReadMode {
40 Insert,
42 Delete {
44 parent_deletes: Vec<FileScanTaskDeleteFile>,
45 },
46 Remove,
48}
49
50pub 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(¤t),
62 "deleted positions shrank for a retained data file"
63 );
64 current -= parent;
65 Ok(current)
66}
67
68#[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 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;