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    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    metrics: Option<BatchMetrics>,
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        Self {
75            iceberg_config,
76            chunk_size,
77            schema,
78            file_scan_tasks: Some(file_scan_tasks),
79            identity,
80            metrics,
81            need_seq_num,
82            need_file_path_and_pos,
83            limit,
84        }
85    }
86
87    #[try_stream(ok = DataChunk, error = BatchError)]
88    async fn do_execute(mut self: Box<Self>) {
89        let table = self.iceberg_config.load_table().await?;
90        let data_types = self.schema.data_types();
91
92        let data_file_scan_tasks = match Option::take(&mut self.file_scan_tasks) {
93            Some(file_scan_tasks) => file_scan_tasks.into_tasks(),
94            None => {
95                bail!("file_scan_tasks must be Some")
96            }
97        };
98        let mut remaining_limit = self
99            .limit
100            .map(|limit| usize::try_from(limit).unwrap_or(usize::MAX));
101
102        for data_file_scan_task in data_file_scan_tasks {
103            if matches!(remaining_limit, Some(0)) {
104                return Ok(());
105            }
106
107            #[for_await]
108            for chunk in scan_task_to_chunk_with_deletes(
109                table.clone(),
110                data_file_scan_task,
111                IcebergScanOpts {
112                    chunk_size: self.chunk_size,
113                    need_seq_num: self.need_seq_num,
114                    need_file_path_and_pos: self.need_file_path_and_pos,
115                    // Iceberg V2 scans expose delete files separately for delete-file scan nodes.
116                    // From V3 onward, deletion vectors are attached to data-file tasks and must be
117                    // applied by iceberg-rs during the data scan.
118                    handle_delete_files: table.metadata().format_version()
119                        >= iceberg::spec::FormatVersion::V3,
120                },
121                self.metrics.as_ref().map(|m| m.iceberg_scan_metrics()),
122            ) {
123                let chunk = chunk?;
124                assert_eq!(chunk.data_types(), data_types);
125                if let Some(remaining) = &mut remaining_limit {
126                    if chunk.cardinality() > *remaining {
127                        yield take_first_visible_rows(chunk, *remaining);
128                        return Ok(());
129                    }
130
131                    *remaining -= chunk.cardinality();
132                    yield chunk;
133
134                    if *remaining == 0 {
135                        return Ok(());
136                    }
137                } else {
138                    yield chunk;
139                }
140            }
141        }
142    }
143}
144
145pub struct IcebergScanExecutorBuilder {}
146
147impl BoxedExecutorBuilder for IcebergScanExecutorBuilder {
148    async fn new_boxed_executor(
149        source: &ExecutorBuilder<'_>,
150        inputs: Vec<BoxedExecutor>,
151    ) -> crate::error::Result<BoxedExecutor> {
152        ensure!(
153            inputs.is_empty(),
154            "Iceberg source should not have input executor!"
155        );
156        let source_node = try_match_expand!(
157            source.plan_node().get_node_body().unwrap(),
158            NodeBody::IcebergScan
159        )?;
160
161        // prepare connector source
162        let options_with_secret = WithOptionsSecResolved::new(
163            source_node.with_properties.clone(),
164            source_node.secret_refs.clone(),
165        );
166        let config = ConnectorProperties::extract(options_with_secret, false)?;
167
168        let split_list = source_node
169            .split
170            .iter()
171            .map(|split| SplitImpl::restore_from_bytes(split).unwrap())
172            .collect_vec();
173        assert_eq!(split_list.len(), 1);
174
175        let fields = source_node
176            .columns
177            .iter()
178            .map(|prost| {
179                let column_desc = prost.column_desc.as_ref().unwrap();
180                let data_type = DataType::from(column_desc.column_type.as_ref().unwrap());
181                let name = column_desc.name.clone();
182                Field::with_name(data_type, name)
183            })
184            .collect();
185        let schema = Schema::new(fields);
186        let metrics = source.context().batch_metrics();
187
188        if let ConnectorProperties::Iceberg(iceberg_properties) = config
189            && let SplitImpl::Iceberg(split) = &split_list[0]
190        {
191            let iceberg_properties: IcebergProperties = *iceberg_properties;
192            let split: IcebergSplit = split.clone();
193            let need_seq_num = schema
194                .fields()
195                .iter()
196                .any(|f| f.name == ICEBERG_SEQUENCE_NUM_COLUMN_NAME);
197            let need_file_path_and_pos = schema
198                .fields()
199                .iter()
200                .any(|f| f.name == ICEBERG_FILE_PATH_COLUMN_NAME)
201                && matches!(split.task, IcebergFileScanTask::Data(_));
202
203            Ok(Box::new(IcebergScanExecutor::new(
204                iceberg_properties,
205                split.task,
206                source.context().get_config().developer.chunk_size,
207                schema,
208                source.plan_node().get_identity().clone(),
209                metrics,
210                need_seq_num,
211                need_file_path_and_pos,
212                split.limit,
213            )))
214        } else {
215            unreachable!()
216        }
217    }
218}
219
220fn take_first_visible_rows(chunk: DataChunk, limit: usize) -> DataChunk {
221    if limit >= chunk.cardinality() {
222        return chunk;
223    }
224
225    let indexes = chunk.visibility().iter_ones().take(limit).collect_vec();
226    chunk.reorder_rows(&indexes)
227}
risingwave_batch_executors/executor/iceberg_scan.rs

risingwave_batch_executors/executor/
iceberg_scan.rs
