risingwave_batch_executors/executor/
iceberg_scan.rs

1// Copyright 2024 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 futures_async_stream::try_stream;
16use futures_util::stream::StreamExt;
17use itertools::Itertools;
18use risingwave_common::array::DataChunk;
19use risingwave_common::catalog::{
20    Field, ICEBERG_FILE_PATH_COLUMN_NAME, ICEBERG_SEQUENCE_NUM_COLUMN_NAME, Schema,
21};
22use risingwave_common::types::DataType;
23use risingwave_connector::WithOptionsSecResolved;
24use risingwave_connector::source::iceberg::{
25    IcebergFileScanMetrics, IcebergFileScanTask, IcebergProperties, IcebergScanOpts, IcebergSplit,
26    scan_task_to_chunk_with_deletes,
27};
28use risingwave_connector::source::{ConnectorProperties, SplitImpl, SplitMetaData};
29use risingwave_pb::batch_plan::plan_node::NodeBody;
30
31use super::{BoxedExecutor, BoxedExecutorBuilder, ExecutorBuilder};
32use crate::error::BatchError;
33use crate::executor::Executor;
34use crate::monitor::BatchMetrics;
35
36pub struct IcebergScanExecutor {
37    iceberg_config: IcebergProperties,
38    file_scan_tasks: Option<IcebergFileScanTask>,
39    chunk_size: usize,
40    schema: Schema,
41    identity: String,
42    file_scan_metrics: Option<IcebergFileScanMetrics>,
43    need_seq_num: bool,
44    need_file_path_and_pos: bool,
45    limit: Option<u64>,
46}
47
48impl Executor for IcebergScanExecutor {
49    fn schema(&self) -> &risingwave_common::catalog::Schema {
50        &self.schema
51    }
52
53    fn identity(&self) -> &str {
54        &self.identity
55    }
56
57    fn execute(self: Box<Self>) -> super::BoxedDataChunkStream {
58        self.do_execute().boxed()
59    }
60}
61
62impl IcebergScanExecutor {
63    pub fn new(
64        iceberg_config: IcebergProperties,
65        file_scan_tasks: IcebergFileScanTask,
66        chunk_size: usize,
67        schema: Schema,
68        identity: String,
69        metrics: Option<BatchMetrics>,
70        need_seq_num: bool,
71        need_file_path_and_pos: bool,
72        limit: Option<u64>,
73    ) -> Self {
74        let file_scan_metrics = metrics.as_ref().map(|metrics| {
75            IcebergFileScanMetrics::new(
76                &metrics.iceberg_scan_metrics(),
77                iceberg_config.table.table_name(),
78            )
79        });
80        Self {
81            iceberg_config,
82            chunk_size,
83            schema,
84            file_scan_tasks: Some(file_scan_tasks),
85            identity,
86            file_scan_metrics,
87            need_seq_num,
88            need_file_path_and_pos,
89            limit,
90        }
91    }
92
93    #[try_stream(ok = DataChunk, error = BatchError)]
94    async fn do_execute(mut self: Box<Self>) {
95        let table = self.iceberg_config.load_table().await?;
96        let data_types = self.schema.data_types();
97
98        let data_file_scan_tasks = match Option::take(&mut self.file_scan_tasks) {
99            Some(file_scan_tasks) => file_scan_tasks.into_tasks(),
100            None => {
101                bail!("file_scan_tasks must be Some")
102            }
103        };
104        let mut remaining_limit = self
105            .limit
106            .map(|limit| usize::try_from(limit).unwrap_or(usize::MAX));
107
108        for data_file_scan_task in data_file_scan_tasks {
109            if matches!(remaining_limit, Some(0)) {
110                return Ok(());
111            }
112
113            #[for_await]
114            for chunk in scan_task_to_chunk_with_deletes(
115                table.clone(),
116                data_file_scan_task,
117                IcebergScanOpts {
118                    chunk_size: self.chunk_size,
119                    need_seq_num: self.need_seq_num,
120                    need_file_path_and_pos: self.need_file_path_and_pos,
121                    // Iceberg V2 scans expose delete files separately for delete-file scan nodes.
122                    // From V3 onward, deletion vectors are attached to data-file tasks and must be
123                    // applied by iceberg-rs during the data scan.
124                    handle_delete_files: table.metadata().format_version()
125                        >= iceberg::spec::FormatVersion::V3,
126                },
127                self.file_scan_metrics.clone(),
128            ) {
129                let chunk = chunk?;
130                assert_eq!(chunk.data_types(), data_types);
131                if let Some(remaining) = &mut remaining_limit {
132                    if chunk.cardinality() > *remaining {
133                        yield take_first_visible_rows(chunk, *remaining);
134                        return Ok(());
135                    }
136
137                    *remaining -= chunk.cardinality();
138                    yield chunk;
139
140                    if *remaining == 0 {
141                        return Ok(());
142                    }
143                } else {
144                    yield chunk;
145                }
146            }
147        }
148    }
149}
150
151pub struct IcebergScanExecutorBuilder {}
152
153impl BoxedExecutorBuilder for IcebergScanExecutorBuilder {
154    async fn new_boxed_executor(
155        source: &ExecutorBuilder<'_>,
156        inputs: Vec<BoxedExecutor>,
157    ) -> crate::error::Result<BoxedExecutor> {
158        ensure!(
159            inputs.is_empty(),
160            "Iceberg source should not have input executor!"
161        );
162        let source_node = try_match_expand!(
163            source.plan_node().get_node_body().unwrap(),
164            NodeBody::IcebergScan
165        )?;
166
167        // prepare connector source
168        let options_with_secret = WithOptionsSecResolved::new(
169            source_node.with_properties.clone(),
170            source_node.secret_refs.clone(),
171        );
172        let config = ConnectorProperties::extract(options_with_secret, false)?;
173
174        let split_list = source_node
175            .split
176            .iter()
177            .map(|split| SplitImpl::restore_from_bytes(split).unwrap())
178            .collect_vec();
179        assert_eq!(split_list.len(), 1);
180
181        let fields = source_node
182            .columns
183            .iter()
184            .map(|prost| {
185                let column_desc = prost.column_desc.as_ref().unwrap();
186                let data_type = DataType::from(column_desc.column_type.as_ref().unwrap());
187                let name = column_desc.name.clone();
188                Field::with_name(data_type, name)
189            })
190            .collect();
191        let schema = Schema::new(fields);
192        let metrics = source.context().batch_metrics();
193
194        if let ConnectorProperties::Iceberg(iceberg_properties) = config
195            && let SplitImpl::Iceberg(split) = &split_list[0]
196        {
197            let iceberg_properties: IcebergProperties = *iceberg_properties;
198            let split: IcebergSplit = split.clone();
199            let need_seq_num = schema
200                .fields()
201                .iter()
202                .any(|f| f.name == ICEBERG_SEQUENCE_NUM_COLUMN_NAME);
203            let need_file_path_and_pos = schema
204                .fields()
205                .iter()
206                .any(|f| f.name == ICEBERG_FILE_PATH_COLUMN_NAME)
207                && matches!(split.task, IcebergFileScanTask::Data(_));
208
209            Ok(Box::new(IcebergScanExecutor::new(
210                iceberg_properties,
211                split.task,
212                source.context().get_config().developer.chunk_size,
213                schema,
214                source.plan_node().get_identity().clone(),
215                metrics,
216                need_seq_num,
217                need_file_path_and_pos,
218                split.limit,
219            )))
220        } else {
221            unreachable!()
222        }
223    }
224}
225
226fn take_first_visible_rows(chunk: DataChunk, limit: usize) -> DataChunk {
227    if limit >= chunk.cardinality() {
228        return chunk;
229    }
230
231    let indexes = chunk.visibility().iter_ones().take(limit).collect_vec();
232    chunk.reorder_rows(&indexes)
233}
risingwave_batch_executors/executor/iceberg_scan.rs

risingwave_batch_executors/executor/
iceberg_scan.rs
