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

risingwave_connector/source/iceberg/
mod.rs
