risingwave_connector/source/iceberg/
mod.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
15pub mod parquet_file_handler;
16pub mod planner;
17
18pub mod metrics;
19use std::collections::{HashMap, HashSet};
20use std::sync::Arc;
21
22use anyhow::anyhow;
23use async_trait::async_trait;
24use futures::StreamExt;
25use futures_async_stream::{for_await, try_stream};
26use iceberg::Catalog;
27use iceberg::expr::{BoundPredicate, Predicate as IcebergPredicate};
28use iceberg::scan::FileScanTask;
29use iceberg::spec::{FormatVersion, TableMetadata};
30use iceberg::table::Table;
31pub use parquet_file_handler::*;
32use phf::{Set, phf_set};
33pub use planner::{
34    IcebergIncrementalScan, IcebergScanMetricsLabels, IcebergScanPlan, IcebergScanPlanner,
35    IcebergScanProjection, IcebergScanTaskBatchMode, IcebergScanTaskPlanner, PersistedFileScanTask,
36};
37use risingwave_common::array::arrow::IcebergArrowConvert;
38use risingwave_common::array::{
39    ArrayBuilder, ArrayImpl, DataChunk, I64Array, Utf8Array, VariantArrayBuilder,
40};
41use risingwave_common::bail;
42use risingwave_common::types::JsonbVal;
43use risingwave_common_estimate_size::EstimateSize;
44use risingwave_pb::batch_plan::iceberg_scan_node::IcebergScanType;
45use serde::{Deserialize, Serialize};
46
47pub use self::metrics::{GLOBAL_ICEBERG_SCAN_METRICS, IcebergFileScanMetrics, IcebergScanMetrics};
48use crate::connector_common::{
49    IcebergCommon, IcebergTableIdentifier, iceberg_java_catalog_props_from_options,
50};
51use crate::enforce_secret::{EnforceSecret, EnforceSecretError};
52use crate::error::{ConnectorError, ConnectorResult};
53use crate::parser::ParserConfig;
54use crate::source::{
55    BoxSourceChunkStream, Column, SourceContextRef, SourceEnumeratorContextRef, SourceProperties,
56    SplitEnumerator, SplitId, SplitMetaData, SplitReader, UnknownFields,
57};
58pub const ICEBERG_CONNECTOR: &str = "iceberg";
59
60#[derive(Clone, Debug, Deserialize, with_options::WithOptions)]
61pub struct IcebergProperties {
62    #[serde(flatten)]
63    pub common: IcebergCommon,
64
65    #[serde(flatten)]
66    pub table: IcebergTableIdentifier,
67
68    // For jdbc catalog
69    #[serde(rename = "catalog.jdbc.user")]
70    pub jdbc_user: Option<String>,
71    #[serde(rename = "catalog.jdbc.password")]
72    pub jdbc_password: Option<String>,
73
74    #[serde(flatten)]
75    pub unknown_fields: HashMap<String, String>,
76}
77
78impl EnforceSecret for IcebergProperties {
79    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
80        "catalog.jdbc.password",
81    };
82
83    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
84        for prop in prop_iter {
85            IcebergCommon::enforce_one(prop)?;
86            if Self::ENFORCE_SECRET_PROPERTIES.contains(prop) {
87                return Err(EnforceSecretError {
88                    key: prop.to_owned(),
89                }
90                .into());
91            }
92        }
93        Ok(())
94    }
95}
96
97impl IcebergProperties {
98    fn java_catalog_props(&self) -> HashMap<String, String> {
99        let mut java_catalog_props = iceberg_java_catalog_props_from_options(
100            self.unknown_fields
101                .iter()
102                .map(|(key, value)| (key.as_str(), value.as_str())),
103        );
104        if let Some(jdbc_user) = self.jdbc_user.clone() {
105            java_catalog_props.insert("jdbc.user".to_owned(), jdbc_user);
106        }
107        if let Some(jdbc_password) = self.jdbc_password.clone() {
108            java_catalog_props.insert("jdbc.password".to_owned(), jdbc_password);
109        }
110        java_catalog_props
111    }
112
113    pub async fn create_catalog(&self) -> ConnectorResult<Arc<dyn Catalog>> {
114        self.common
115            .resolve_catalog_config(self.java_catalog_props())?
116            .create_catalog()
117            .await
118    }
119
120    pub async fn load_table(&self) -> ConnectorResult<Table> {
121        self.common
122            .resolve_catalog_config(self.java_catalog_props())?
123            .load_table(&self.table)
124            .await
125    }
126}
127
128impl SourceProperties for IcebergProperties {
129    type Split = IcebergSplit;
130    type SplitEnumerator = IcebergSplitEnumerator;
131    type SplitReader = IcebergFileReader;
132
133    const SOURCE_NAME: &'static str = ICEBERG_CONNECTOR;
134}
135
136impl UnknownFields for IcebergProperties {
137    fn unknown_fields(&self) -> HashMap<String, String> {
138        self.unknown_fields.clone()
139    }
140}
141
142#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
143pub enum IcebergFileScanTask {
144    Data(Vec<FileScanTask>),
145    EqualityDelete(Vec<FileScanTask>),
146    PositionDelete(Vec<FileScanTask>),
147}
148
149impl IcebergFileScanTask {
150    pub fn tasks(&self) -> &[FileScanTask] {
151        match self {
152            IcebergFileScanTask::Data(file_scan_tasks)
153            | IcebergFileScanTask::EqualityDelete(file_scan_tasks)
154            | IcebergFileScanTask::PositionDelete(file_scan_tasks) => file_scan_tasks,
155        }
156    }
157
158    pub fn is_empty(&self) -> bool {
159        self.tasks().is_empty()
160    }
161
162    pub fn files(&self) -> Vec<String> {
163        self.tasks()
164            .iter()
165            .map(|task| task.data_file_path.clone())
166            .collect()
167    }
168
169    pub fn predicate(&self) -> Option<&BoundPredicate> {
170        let first_task = self.tasks().first()?;
171        first_task.predicate.as_ref()
172    }
173
174    fn strip_non_serializable_planning_context(&mut self) {
175        let tasks = match self {
176            IcebergFileScanTask::Data(tasks)
177            | IcebergFileScanTask::EqualityDelete(tasks)
178            | IcebergFileScanTask::PositionDelete(tasks) => tasks,
179        };
180
181        for task in tasks {
182            // These fields are deliberately not serializable in Iceberg. The previous
183            // Iceberg dependency skipped them during serde, so retain that split format.
184            task.partition = None;
185            task.partition_spec = None;
186            task.name_mapping = None;
187            task.unified_partition_type = None;
188        }
189    }
190}
191
192#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
193pub struct IcebergSplit {
194    pub split_id: i64,
195    pub task: IcebergFileScanTask,
196    #[serde(default, skip_serializing_if = "Option::is_none")]
197    pub limit: Option<u64>,
198}
199
200impl IcebergSplit {
201    #[allow(deprecated)]
202    pub fn empty(iceberg_scan_type: IcebergScanType) -> Self {
203        let task = match iceberg_scan_type {
204            IcebergScanType::DataScan => IcebergFileScanTask::Data(vec![]),
205            IcebergScanType::EqualityDeleteScan => IcebergFileScanTask::EqualityDelete(vec![]),
206            IcebergScanType::PositionDeleteScan => IcebergFileScanTask::PositionDelete(vec![]),
207            IcebergScanType::Unspecified | IcebergScanType::CountStar => {
208                // These scan types do not carry file tasks. Keep the split serializable without
209                // introducing a new empty-task variant.
210                IcebergFileScanTask::Data(vec![])
211            }
212        };
213        Self {
214            split_id: 0,
215            task,
216            limit: None,
217        }
218    }
219}
220
221impl SplitMetaData for IcebergSplit {
222    fn id(&self) -> SplitId {
223        self.split_id.to_string().into()
224    }
225
226    fn restore_from_json(value: JsonbVal) -> ConnectorResult<Self> {
227        serde_json::from_value(value.take()).map_err(|e| anyhow!(e).into())
228    }
229
230    fn encode_to_json(&self) -> JsonbVal {
231        let mut split = self.clone();
232        split.task.strip_non_serializable_planning_context();
233        serde_json::to_value(split)
234            .expect("iceberg split serialization should not fail")
235            .into()
236    }
237
238    fn update_offset(&mut self, _last_seen_offset: String) -> ConnectorResult<()> {
239        // Iceberg source progress is tracked by persisted file tasks in the stream state table.
240        // A split does not carry an intra-file offset until partial-file reads are supported.
241        Ok(())
242    }
243}
244
245#[derive(Debug, Clone)]
246pub struct IcebergSplitEnumerator {
247    config: IcebergProperties,
248}
249
250#[derive(Debug, Clone)]
251pub struct IcebergDeleteParameters {
252    pub equality_delete_columns: Vec<String>,
253    pub has_position_delete: bool,
254    pub snapshot_id: Option<i64>,
255}
256
257#[async_trait]
258impl SplitEnumerator for IcebergSplitEnumerator {
259    type Properties = IcebergProperties;
260    type Split = IcebergSplit;
261
262    async fn new(
263        properties: Self::Properties,
264        context: SourceEnumeratorContextRef,
265    ) -> ConnectorResult<Self> {
266        Ok(Self::new_inner(properties, context))
267    }
268
269    async fn list_splits(&mut self) -> ConnectorResult<Vec<Self::Split>> {
270        // Like file source, iceberg streaming source has a List Executor and a Fetch Executor,
271        // instead of relying on SplitEnumerator on meta.
272        // TODO: add some validation logic here.
273        Ok(vec![])
274    }
275}
276impl IcebergSplitEnumerator {
277    pub fn new_inner(properties: IcebergProperties, _context: SourceEnumeratorContextRef) -> Self {
278        Self { config: properties }
279    }
280}
281
282#[derive(Debug, Clone, PartialEq, Eq, Hash)]
283pub enum IcebergTimeTravelInfo {
284    Version(i64),
285    TimestampMs(i64),
286}
287
288#[derive(Debug, Clone)]
289pub struct IcebergListResult {
290    pub data_files: Vec<FileScanTask>,
291    pub equality_delete_files: Vec<FileScanTask>,
292    pub position_delete_files: Vec<FileScanTask>,
293    pub equality_delete_columns: Vec<String>,
294    pub format_version: FormatVersion,
295    pub schema: std::sync::Arc<iceberg::spec::Schema>,
296}
297
298impl IcebergSplitEnumerator {
299    pub fn get_snapshot_id(
300        table: &Table,
301        time_travel_info: Option<IcebergTimeTravelInfo>,
302    ) -> ConnectorResult<Option<i64>> {
303        Self::get_snapshot_id_from_metadata(table.metadata(), time_travel_info)
304    }
305
306    fn get_snapshot_id_from_metadata(
307        metadata: &TableMetadata,
308        time_travel_info: Option<IcebergTimeTravelInfo>,
309    ) -> ConnectorResult<Option<i64>> {
310        let snapshot_id = match time_travel_info {
311            Some(IcebergTimeTravelInfo::Version(version)) => {
312                let Some(snapshot) = metadata.snapshot_by_id(version) else {
313                    bail!("Cannot find the snapshot id in the iceberg table.");
314                };
315                Some(snapshot.snapshot_id())
316            }
317            Some(IcebergTimeTravelInfo::TimestampMs(timestamp)) => {
318                let snapshot_log = metadata
319                    .history()
320                    .iter()
321                    .rev()
322                    .find(|snapshot_log| snapshot_log.timestamp_ms() <= timestamp);
323                match snapshot_log {
324                    Some(snapshot_log) => Some(snapshot_log.snapshot_id),
325                    None => {
326                        // convert unix time to human-readable time
327                        let time = chrono::DateTime::from_timestamp_millis(timestamp);
328                        if let Some(time) = time {
329                            tracing::warn!("Cannot find a snapshot older than {}", time);
330                        } else {
331                            tracing::warn!("Cannot find a snapshot");
332                        }
333                        return Ok(None);
334                    }
335                }
336            }
337            None => metadata.current_snapshot_id(),
338        };
339        Ok(snapshot_id)
340    }
341
342    pub async fn list_scan_tasks(
343        &self,
344        time_travel_info: Option<IcebergTimeTravelInfo>,
345        predicate: IcebergPredicate,
346    ) -> ConnectorResult<Option<IcebergListResult>> {
347        let table = self.config.load_table().await?;
348        let snapshot_id = Self::get_snapshot_id(&table, time_travel_info)?;
349
350        let Some(snapshot_id) = snapshot_id else {
351            return Ok(None);
352        };
353        let res = self
354            .list_scan_tasks_inner(&table, snapshot_id, predicate)
355            .await?;
356        Ok(Some(res))
357    }
358
359    async fn list_scan_tasks_inner(
360        &self,
361        table: &Table,
362        snapshot_id: i64,
363        predicate: IcebergPredicate,
364    ) -> ConnectorResult<IcebergListResult> {
365        let format_version = table.metadata().format_version();
366        let table_schema = table.metadata().current_schema();
367        tracing::debug!("iceberg_table_schema: {:?}", table_schema);
368
369        let mut position_delete_files = vec![];
370        let mut position_delete_files_set = HashSet::new();
371        let mut data_files = vec![];
372        let mut equality_delete_files = vec![];
373        let mut equality_delete_files_set = HashSet::new();
374        let mut equality_delete_ids = None;
375        let mut scan_builder = table.scan().snapshot_id(snapshot_id).select_all();
376        if predicate != IcebergPredicate::AlwaysTrue {
377            scan_builder = scan_builder.with_filter(predicate.clone());
378        }
379        let scan = scan_builder.build()?;
380        let file_scan_stream = scan.plan_files().await?;
381
382        #[for_await]
383        for task in file_scan_stream {
384            let task: FileScanTask = task?;
385
386            // Collect delete files for separate scan types, but keep task.deletes intact
387            for delete_file in &task.deletes {
388                match delete_file.file_type {
389                    iceberg::spec::DataContentType::Data => {
390                        bail!("Data file should not in task deletes");
391                    }
392                    iceberg::spec::DataContentType::EqualityDeletes => {
393                        if equality_delete_files_set.insert(delete_file.file_path.clone()) {
394                            if equality_delete_ids.is_none() {
395                                equality_delete_ids = delete_file.equality_ids.clone();
396                            } else if equality_delete_ids != delete_file.equality_ids {
397                                bail!(
398                                    "The schema of iceberg equality delete file must be consistent"
399                                );
400                            }
401                            equality_delete_files.push(delete_file.to_file_scan_task(&task));
402                        }
403                    }
404                    iceberg::spec::DataContentType::PositionDeletes => {
405                        if position_delete_files_set.insert(delete_file.file_path.clone()) {
406                            position_delete_files.push(delete_file.to_file_scan_task(&task));
407                        }
408                    }
409                }
410            }
411
412            // Top-level scan tasks always represent data files. Keep their delete
413            // descriptors intact so the SDK reader can apply them when requested.
414            data_files.push(task);
415        }
416        let schema = table_schema.clone();
417        let equality_delete_columns = equality_delete_ids
418            .unwrap_or_default()
419            .into_iter()
420            .map(|id| match schema.name_by_field_id(id) {
421                Some(name) => Ok::<std::string::String, ConnectorError>(name.to_owned()),
422                None => bail!("Delete field id {} not found in schema", id),
423            })
424            .collect::<ConnectorResult<Vec<_>>>()?;
425
426        Ok(IcebergListResult {
427            data_files,
428            equality_delete_files,
429            position_delete_files,
430            equality_delete_columns,
431            format_version,
432            schema,
433        })
434    }
435
436    /// Uniformly distribute scan tasks to compute nodes.
437    /// It's deterministic so that it can best utilize the data locality.
438    ///
439    /// # Arguments
440    /// * `file_scan_tasks`: The file scan tasks to be split.
441    /// * `split_num`: The number of splits to be created.
442    ///
443    /// This algorithm is based on a min-heap. It will push all groups into the heap, and then pop the smallest group and add the file scan task to it.
444    /// Ensure that the total length of each group is as balanced as possible.
445    /// The time complexity is O(n log k), where n is the number of file scan tasks and k is the number of splits.
446    /// The space complexity is O(k), where k is the number of splits.
447    /// The algorithm is stable, so the order of the file scan tasks will be preserved.
448    pub fn split_n_vecs(
449        file_scan_tasks: Vec<FileScanTask>,
450        split_num: usize,
451    ) -> Vec<Vec<FileScanTask>> {
452        IcebergScanTaskPlanner::split_n_vecs(file_scan_tasks, split_num)
453    }
454}
455
456pub struct IcebergScanOpts {
457    pub chunk_size: usize,
458    pub need_seq_num: bool,
459    pub need_file_path_and_pos: bool,
460    pub handle_delete_files: bool,
461}
462
463/// Scan a data file. Delete files are handled by the iceberg-rust `reader.read` implementation.
464#[try_stream(ok = DataChunk, error = ConnectorError)]
465pub async fn scan_task_to_chunk_with_deletes(
466    table: Table,
467    mut data_file_scan_task: FileScanTask,
468    IcebergScanOpts {
469        chunk_size,
470        need_seq_num,
471        need_file_path_and_pos,
472        handle_delete_files,
473    }: IcebergScanOpts,
474    metrics: Option<IcebergFileScanMetrics>,
475) {
476    let num_delete_files = data_file_scan_task.deletes.len();
477    let expected_record_count = data_file_scan_task.record_count;
478    let file_start = std::time::Instant::now();
479
480    let read_metrics = metrics.clone();
481    let mut read_bytes = scopeguard::guard(0u64, move |read_bytes| {
482        if let Some(metrics) = read_metrics {
483            metrics.record_read_bytes(read_bytes);
484        }
485    });
486
487    let data_file_path = data_file_scan_task.data_file_path.clone();
488    let data_sequence_number = data_file_scan_task.sequence_number;
489
490    tracing::debug!(
491        "scan_task_to_chunk_with_deletes: data_file={}, handle_delete_files={}, total_delete_files={}",
492        data_file_path,
493        handle_delete_files,
494        data_file_scan_task.deletes.len()
495    );
496
497    if !handle_delete_files {
498        // Keep the delete files from being applied when the caller opts out.
499        data_file_scan_task.deletes.clear();
500    }
501
502    // Read the data file; delete application is delegated to the reader.
503    let reader = table
504        .reader_builder()
505        .with_batch_size(chunk_size)
506        .with_row_group_filtering_enabled(true)
507        .build();
508    let file_scan_stream = tokio_stream::once(Ok(data_file_scan_task.clone()));
509
510    let mut record_batch_stream: iceberg::scan::ArrowRecordBatchStream =
511        reader.read(Box::pin(file_scan_stream))?.stream();
512
513    // The reader rejects a file with shredded variant columns before yielding any batch.
514    // Retry without the variant columns; NULL columns are spliced back in below.
515    let mut null_padded_variant_positions: Option<Vec<usize>> = None;
516    let record_batch_stream: iceberg::scan::ArrowRecordBatchStream =
517        match record_batch_stream.next().await {
518            Some(Err(e)) if is_shredded_variant_rejection(&e) => {
519                let (variant_field_ids, variant_positions, variant_names) =
520                    projected_variant_columns(&data_file_scan_task);
521                if variant_field_ids.is_empty() {
522                    return Err(e.into());
523                }
524                tracing::warn!(
525                    data_file_path,
526                    columns = ?variant_names,
527                    "shredded variant columns are not supported yet; reading them as NULL",
528                );
529                null_padded_variant_positions = Some(variant_positions);
530
531                let mut reduced_task = data_file_scan_task;
532                reduced_task
533                    .project_field_ids
534                    .retain(|id| !variant_field_ids.contains(id));
535                let reader = table
536                    .reader_builder()
537                    .with_batch_size(chunk_size)
538                    .with_row_group_filtering_enabled(true)
539                    .build();
540                reader
541                    .read(Box::pin(tokio_stream::once(Ok(reduced_task))))?
542                    .stream()
543            }
544            first => Box::pin(futures::stream::iter(first).chain(record_batch_stream)),
545        };
546    let mut record_batch_stream = record_batch_stream.enumerate();
547
548    let mut total_rows_read: u64 = 0;
549
550    // Process each record batch. Delete application is handled by the SDK.
551    while let Some((batch_index, record_batch)) = record_batch_stream.next().await {
552        let record_batch = record_batch?;
553        let batch_start_pos = (batch_index * chunk_size) as i64;
554
555        let mut chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
556        if let Some(positions) = &null_padded_variant_positions {
557            chunk = pad_null_variant_columns(chunk, positions, record_batch.num_rows());
558        }
559        let row_count = chunk.capacity();
560        total_rows_read += row_count as u64;
561
562        // Add metadata columns if requested
563        if need_seq_num {
564            let (mut columns, visibility) = chunk.into_parts();
565            columns.push(Arc::new(ArrayImpl::Int64(I64Array::from_iter(
566                std::iter::repeat_n(data_sequence_number, row_count),
567            ))));
568            chunk = DataChunk::from_parts(columns.into(), visibility);
569        }
570
571        if need_file_path_and_pos {
572            let (mut columns, visibility) = chunk.into_parts();
573            columns.push(Arc::new(ArrayImpl::Utf8(Utf8Array::from_iter(
574                std::iter::repeat_n(data_file_path.as_str(), row_count),
575            ))));
576
577            // Generate position values for each row in the batch
578            let positions: Vec<i64> =
579                (batch_start_pos..(batch_start_pos + row_count as i64)).collect();
580            columns.push(Arc::new(ArrayImpl::Int64(I64Array::from_iter(positions))));
581
582            chunk = DataChunk::from_parts(columns.into(), visibility);
583        }
584
585        *read_bytes += chunk.estimated_heap_size() as u64;
586        yield chunk;
587    }
588
589    // Record per-file metrics after reading all batches.
590    if let Some(metrics) = metrics {
591        metrics.record_file_read_duration(file_start.elapsed().as_secs_f64());
592
593        if total_rows_read > 0 {
594            metrics.record_rows_read(total_rows_read);
595        }
596
597        metrics.record_file_read();
598
599        // APPROXIMATE: Estimate delete rows applied. The delta between expected_record_count
600        // and actual rows read may also include predicate pushdown / row-group pruning effects,
601        // so this metric can overcount. It is still useful as an approximate signal for
602        // detecting whether delete files cause significant row filtering.
603        if handle_delete_files
604            && num_delete_files > 0
605            && let Some(expected) = expected_record_count
606        {
607            let deleted = expected.saturating_sub(total_rows_read);
608            if deleted > 0 {
609                metrics.record_delete_rows_applied(deleted);
610            }
611        }
612    }
613}
614
615/// Whether the error is the reader's per-file rejection of shredded variant columns.
616// `IcebergError` hides the inner error, so the raw one is needed to inspect kind/message.
617#[expect(clippy::disallowed_types)]
618fn is_shredded_variant_rejection(e: &iceberg::Error) -> bool {
619    e.kind() == iceberg::ErrorKind::FeatureUnsupported && e.message().contains("shredded variant")
620}
621
622/// The projected VARIANT columns of a task: their field ids, their positions in the
623/// projected column order, and their names.
624fn projected_variant_columns(task: &FileScanTask) -> (Vec<i32>, Vec<usize>, Vec<String>) {
625    let mut field_ids = Vec::new();
626    let mut positions = Vec::new();
627    let mut names = Vec::new();
628    for (position, field_id) in task.project_field_ids.iter().enumerate() {
629        if let Some(field) = task.schema.field_by_id(*field_id)
630            && matches!(field.field_type.as_ref(), iceberg::spec::Type::Variant(_))
631        {
632            field_ids.push(*field_id);
633            positions.push(position);
634            names.push(field.name.clone());
635        }
636    }
637    (field_ids, positions, names)
638}
639
640/// Insert all-NULL variant columns at the given projected positions.
641fn pad_null_variant_columns(chunk: DataChunk, positions: &[usize], row_count: usize) -> DataChunk {
642    let (mut columns, visibility) = chunk.into_parts();
643    for &position in positions {
644        let mut builder = VariantArrayBuilder::new(row_count);
645        for _ in 0..row_count {
646            builder.append_null();
647        }
648        columns.insert(position, Arc::new(ArrayImpl::Variant(builder.finish())));
649    }
650    DataChunk::from_parts(columns.into(), visibility)
651}
652
653#[derive(Debug)]
654pub struct IcebergFileReader {}
655
656#[async_trait]
657impl SplitReader for IcebergFileReader {
658    type Properties = IcebergProperties;
659    type Split = IcebergSplit;
660
661    async fn new(
662        _props: IcebergProperties,
663        _splits: Vec<IcebergSplit>,
664        _parser_config: ParserConfig,
665        _source_ctx: SourceContextRef,
666        _columns: Option<Vec<Column>>,
667    ) -> ConnectorResult<Self> {
668        unimplemented!()
669    }
670
671    fn into_stream(self) -> BoxSourceChunkStream {
672        unimplemented!()
673    }
674}
675
676#[cfg(test)]
677mod tests {
678    use std::collections::HashMap;
679    use std::sync::Arc;
680
681    use iceberg::scan::FileScanTask;
682    use iceberg::spec::{
683        FormatVersion, MAIN_BRANCH, NestedField, Operation, PrimitiveType, Schema, Snapshot,
684        SortOrder, Struct, Summary, TableMetadataBuilder, Type, UnboundPartitionSpec,
685    };
686
687    use super::*;
688
689    fn test_snapshot(
690        snapshot_id: i64,
691        parent_snapshot_id: Option<i64>,
692        timestamp_ms: i64,
693    ) -> Snapshot {
694        Snapshot::builder()
695            .with_snapshot_id(snapshot_id)
696            .with_parent_snapshot_id(parent_snapshot_id)
697            .with_sequence_number(snapshot_id)
698            .with_timestamp_ms(timestamp_ms)
699            .with_manifest_list(format!("/snap-{snapshot_id}.avro"))
700            .with_summary(Summary {
701                operation: Operation::Append,
702                additional_properties: HashMap::new(),
703            })
704            .with_schema_id(0)
705            .build()
706    }
707
708    fn test_table_metadata_builder() -> TableMetadataBuilder {
709        TableMetadataBuilder::new(
710            Schema::builder()
711                .with_fields(vec![
712                    NestedField::new(1, "id", Type::Primitive(PrimitiveType::Long), false).into(),
713                ])
714                .build()
715                .unwrap(),
716            UnboundPartitionSpec::builder().build(),
717            SortOrder::unsorted_order(),
718            "s3://warehouse/db/table".to_owned(),
719            FormatVersion::V2,
720            HashMap::new(),
721        )
722        .unwrap()
723    }
724
725    #[test]
726    fn test_get_snapshot_id_uses_main_branch_history_for_timestamp() {
727        let metadata = test_table_metadata_builder()
728            .set_branch_snapshot(test_snapshot(1, None, 1_000), MAIN_BRANCH)
729            .unwrap()
730            .build()
731            .unwrap()
732            .metadata;
733        let metadata = metadata
734            .into_builder(Some("s3://warehouse/db/table/v2.metadata.json".to_owned()))
735            .set_branch_snapshot(test_snapshot(2, Some(1), 2_000), MAIN_BRANCH)
736            .unwrap()
737            .build()
738            .unwrap()
739            .metadata;
740        let metadata = metadata
741            .into_builder(Some("s3://warehouse/db/table/v3.metadata.json".to_owned()))
742            .set_branch_snapshot(test_snapshot(3, Some(1), 3_000), "audit")
743            .unwrap()
744            .build()
745            .unwrap()
746            .metadata;
747
748        assert_eq!(
749            IcebergSplitEnumerator::get_snapshot_id_from_metadata(
750                &metadata,
751                Some(IcebergTimeTravelInfo::TimestampMs(3_500)),
752            )
753            .unwrap(),
754            Some(2)
755        );
756        assert_eq!(
757            IcebergSplitEnumerator::get_snapshot_id_from_metadata(
758                &metadata,
759                Some(IcebergTimeTravelInfo::TimestampMs(1_500)),
760            )
761            .unwrap(),
762            Some(1)
763        );
764    }
765
766    #[test]
767    fn test_get_snapshot_id_version_without_current_snapshot() {
768        let metadata = test_table_metadata_builder()
769            .add_snapshot(test_snapshot(7, None, 1_000))
770            .unwrap()
771            .build()
772            .unwrap()
773            .metadata;
774
775        assert_eq!(metadata.current_snapshot_id(), None);
776        assert_eq!(
777            IcebergSplitEnumerator::get_snapshot_id_from_metadata(
778                &metadata,
779                Some(IcebergTimeTravelInfo::Version(7)),
780            )
781            .unwrap(),
782            Some(7)
783        );
784        assert_eq!(
785            IcebergSplitEnumerator::get_snapshot_id_from_metadata(&metadata, None).unwrap(),
786            None
787        );
788    }
789
790    fn create_file_scan_task(length: u64, id: u64) -> FileScanTask {
791        FileScanTask {
792            length,
793            start: 0,
794            record_count: Some(0),
795            first_row_id: None,
796            data_sequence_number: None,
797            data_file_path: format!("test_{}.parquet", id),
798            data_file_format: iceberg::spec::DataFileFormat::Parquet,
799            schema: Arc::new(Schema::builder().build().unwrap()),
800            project_field_ids: vec![],
801            predicate: None,
802            deletes: vec![],
803            sequence_number: 0,
804            file_size_in_bytes: 0,
805            partition: None,
806            partition_spec: None,
807            name_mapping: None,
808            unified_partition_type: None,
809            case_sensitive: true,
810            key_metadata: None,
811        }
812    }
813
814    #[test]
815    fn test_split_serialization_strips_iceberg_planning_context() {
816        let mut task = create_file_scan_task(100, 1);
817        task.partition = Some(Struct::empty());
818        task.partition_spec = Some(Arc::new(iceberg::spec::PartitionSpec::unpartition_spec()));
819
820        let split = IcebergSplit {
821            split_id: 1,
822            task: IcebergFileScanTask::Data(vec![task]),
823            limit: None,
824        };
825        let restored = IcebergSplit::restore_from_json(split.encode_to_json()).unwrap();
826        let task = &restored.task.tasks()[0];
827
828        assert!(task.partition.is_none());
829        assert!(task.partition_spec.is_none());
830        assert!(task.name_mapping.is_none());
831        assert!(task.unified_partition_type.is_none());
832    }
833
834    #[test]
835    fn test_split_n_vecs_basic() {
836        let file_scan_tasks = (1..=12)
837            .map(|i| create_file_scan_task(i + 100, i))
838            .collect::<Vec<_>>(); // Ensure the correct function is called
839
840        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
841
842        assert_eq!(groups.len(), 3);
843
844        let group_lengths: Vec<u64> = groups
845            .iter()
846            .map(|group| group.iter().map(|task| task.length).sum())
847            .collect();
848
849        let max_length = *group_lengths.iter().max().unwrap();
850        let min_length = *group_lengths.iter().min().unwrap();
851        assert!(max_length - min_length <= 10, "Groups should be balanced");
852
853        let total_tasks: usize = groups.iter().map(|group| group.len()).sum();
854        assert_eq!(total_tasks, 12);
855    }
856
857    #[test]
858    fn test_split_n_vecs_empty() {
859        let file_scan_tasks = Vec::new();
860        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
861        assert_eq!(groups.len(), 3);
862        assert!(groups.iter().all(|group| group.is_empty()));
863    }
864
865    #[test]
866    fn test_split_n_vecs_single_task() {
867        let file_scan_tasks = vec![create_file_scan_task(100, 1)];
868        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
869        assert_eq!(groups.len(), 3);
870        assert_eq!(groups.iter().filter(|group| !group.is_empty()).count(), 1);
871    }
872
873    #[test]
874    fn test_split_n_vecs_uneven_distribution() {
875        let file_scan_tasks = vec![
876            create_file_scan_task(1000, 1),
877            create_file_scan_task(100, 2),
878            create_file_scan_task(100, 3),
879            create_file_scan_task(100, 4),
880            create_file_scan_task(100, 5),
881        ];
882
883        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 2);
884        assert_eq!(groups.len(), 2);
885
886        let group_with_large_task = groups
887            .iter()
888            .find(|group| group.iter().any(|task| task.length == 1000))
889            .unwrap();
890        assert_eq!(group_with_large_task.len(), 1);
891    }
892
893    #[test]
894    fn test_split_n_vecs_same_files_distribution() {
895        let file_scan_tasks = vec![
896            create_file_scan_task(100, 1),
897            create_file_scan_task(100, 2),
898            create_file_scan_task(100, 3),
899            create_file_scan_task(100, 4),
900            create_file_scan_task(100, 5),
901            create_file_scan_task(100, 6),
902            create_file_scan_task(100, 7),
903            create_file_scan_task(100, 8),
904        ];
905
906        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks.clone(), 4)
907            .iter()
908            .map(|g| {
909                g.iter()
910                    .map(|task| task.data_file_path.clone())
911                    .collect::<Vec<_>>()
912            })
913            .collect::<Vec<_>>();
914
915        for _ in 0..10000 {
916            let groups_2 = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks.clone(), 4)
917                .iter()
918                .map(|g| {
919                    g.iter()
920                        .map(|task| task.data_file_path.clone())
921                        .collect::<Vec<_>>()
922                })
923                .collect::<Vec<_>>();
924
925            assert_eq!(groups, groups_2);
926        }
927    }
928}
risingwave_connector/source/iceberg/mod.rs

risingwave_connector/source/iceberg/
mod.rs
