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;
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::{ArrayImpl, DataChunk, I64Array, Utf8Array};
39use risingwave_common::bail;
40use risingwave_common::types::JsonbVal;
41use risingwave_common_estimate_size::EstimateSize;
42use risingwave_pb::batch_plan::iceberg_scan_node::IcebergScanType;
43use serde::{Deserialize, Serialize};
44
45pub use self::metrics::{GLOBAL_ICEBERG_SCAN_METRICS, IcebergScanMetrics};
46use crate::connector_common::{
47    IcebergCommon, IcebergTableIdentifier, iceberg_java_catalog_props_from_options,
48};
49use crate::enforce_secret::{EnforceSecret, EnforceSecretError};
50use crate::error::{ConnectorError, ConnectorResult};
51use crate::parser::ParserConfig;
52use crate::source::{
53    BoxSourceChunkStream, Column, SourceContextRef, SourceEnumeratorContextRef, SourceProperties,
54    SplitEnumerator, SplitId, SplitMetaData, SplitReader, UnknownFields,
55};
56pub const ICEBERG_CONNECTOR: &str = "iceberg";
57
58#[derive(Clone, Debug, Deserialize, with_options::WithOptions)]
59pub struct IcebergProperties {
60    #[serde(flatten)]
61    pub common: IcebergCommon,
62
63    #[serde(flatten)]
64    pub table: IcebergTableIdentifier,
65
66    // For jdbc catalog
67    #[serde(rename = "catalog.jdbc.user")]
68    pub jdbc_user: Option<String>,
69    #[serde(rename = "catalog.jdbc.password")]
70    pub jdbc_password: Option<String>,
71
72    #[serde(flatten)]
73    pub unknown_fields: HashMap<String, String>,
74}
75
76impl EnforceSecret for IcebergProperties {
77    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
78        "catalog.jdbc.password",
79    };
80
81    fn enforce_secret<'a>(prop_iter: impl Iterator<Item = &'a str>) -> ConnectorResult<()> {
82        for prop in prop_iter {
83            IcebergCommon::enforce_one(prop)?;
84            if Self::ENFORCE_SECRET_PROPERTIES.contains(prop) {
85                return Err(EnforceSecretError {
86                    key: prop.to_owned(),
87                }
88                .into());
89            }
90        }
91        Ok(())
92    }
93}
94
95impl IcebergProperties {
96    fn java_catalog_props(&self) -> HashMap<String, String> {
97        let mut java_catalog_props = iceberg_java_catalog_props_from_options(
98            self.unknown_fields
99                .iter()
100                .map(|(key, value)| (key.as_str(), value.as_str())),
101        );
102        if let Some(jdbc_user) = self.jdbc_user.clone() {
103            java_catalog_props.insert("jdbc.user".to_owned(), jdbc_user);
104        }
105        if let Some(jdbc_password) = self.jdbc_password.clone() {
106            java_catalog_props.insert("jdbc.password".to_owned(), jdbc_password);
107        }
108        java_catalog_props
109    }
110
111    pub async fn create_catalog(&self) -> ConnectorResult<Arc<dyn Catalog>> {
112        self.common
113            .resolve_catalog_config(self.java_catalog_props())?
114            .create_catalog()
115            .await
116    }
117
118    pub async fn load_table(&self) -> ConnectorResult<Table> {
119        self.common
120            .resolve_catalog_config(self.java_catalog_props())?
121            .load_table(&self.table)
122            .await
123    }
124}
125
126impl SourceProperties for IcebergProperties {
127    type Split = IcebergSplit;
128    type SplitEnumerator = IcebergSplitEnumerator;
129    type SplitReader = IcebergFileReader;
130
131    const SOURCE_NAME: &'static str = ICEBERG_CONNECTOR;
132}
133
134impl UnknownFields for IcebergProperties {
135    fn unknown_fields(&self) -> HashMap<String, String> {
136        self.unknown_fields.clone()
137    }
138}
139
140#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
141pub enum IcebergFileScanTask {
142    Data(Vec<FileScanTask>),
143    EqualityDelete(Vec<FileScanTask>),
144    PositionDelete(Vec<FileScanTask>),
145}
146
147impl IcebergFileScanTask {
148    pub fn tasks(&self) -> &[FileScanTask] {
149        match self {
150            IcebergFileScanTask::Data(file_scan_tasks)
151            | IcebergFileScanTask::EqualityDelete(file_scan_tasks)
152            | IcebergFileScanTask::PositionDelete(file_scan_tasks) => file_scan_tasks,
153        }
154    }
155
156    pub fn is_empty(&self) -> bool {
157        self.tasks().is_empty()
158    }
159
160    pub fn files(&self) -> Vec<String> {
161        self.tasks()
162            .iter()
163            .map(|task| task.data_file_path.clone())
164            .collect()
165    }
166
167    pub fn predicate(&self) -> Option<&BoundPredicate> {
168        let first_task = self.tasks().first()?;
169        first_task.predicate.as_ref()
170    }
171}
172
173#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
174pub struct IcebergSplit {
175    pub split_id: i64,
176    pub task: IcebergFileScanTask,
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub limit: Option<u64>,
179}
180
181impl IcebergSplit {
182    #[allow(deprecated)]
183    pub fn empty(iceberg_scan_type: IcebergScanType) -> Self {
184        let task = match iceberg_scan_type {
185            IcebergScanType::DataScan => IcebergFileScanTask::Data(vec![]),
186            IcebergScanType::EqualityDeleteScan => IcebergFileScanTask::EqualityDelete(vec![]),
187            IcebergScanType::PositionDeleteScan => IcebergFileScanTask::PositionDelete(vec![]),
188            IcebergScanType::Unspecified | IcebergScanType::CountStar => {
189                // These scan types do not carry file tasks. Keep the split serializable without
190                // introducing a new empty-task variant.
191                IcebergFileScanTask::Data(vec![])
192            }
193        };
194        Self {
195            split_id: 0,
196            task,
197            limit: None,
198        }
199    }
200}
201
202impl SplitMetaData for IcebergSplit {
203    fn id(&self) -> SplitId {
204        self.split_id.to_string().into()
205    }
206
207    fn restore_from_json(value: JsonbVal) -> ConnectorResult<Self> {
208        serde_json::from_value(value.take()).map_err(|e| anyhow!(e).into())
209    }
210
211    fn encode_to_json(&self) -> JsonbVal {
212        serde_json::to_value(self.clone())
213            .expect("iceberg split serialization should not fail")
214            .into()
215    }
216
217    fn update_offset(&mut self, _last_seen_offset: String) -> ConnectorResult<()> {
218        // Iceberg source progress is tracked by persisted file tasks in the stream state table.
219        // A split does not carry an intra-file offset until partial-file reads are supported.
220        Ok(())
221    }
222}
223
224#[derive(Debug, Clone)]
225pub struct IcebergSplitEnumerator {
226    config: IcebergProperties,
227}
228
229#[derive(Debug, Clone)]
230pub struct IcebergDeleteParameters {
231    pub equality_delete_columns: Vec<String>,
232    pub has_position_delete: bool,
233    pub snapshot_id: Option<i64>,
234}
235
236#[async_trait]
237impl SplitEnumerator for IcebergSplitEnumerator {
238    type Properties = IcebergProperties;
239    type Split = IcebergSplit;
240
241    async fn new(
242        properties: Self::Properties,
243        context: SourceEnumeratorContextRef,
244    ) -> ConnectorResult<Self> {
245        Ok(Self::new_inner(properties, context))
246    }
247
248    async fn list_splits(&mut self) -> ConnectorResult<Vec<Self::Split>> {
249        // Like file source, iceberg streaming source has a List Executor and a Fetch Executor,
250        // instead of relying on SplitEnumerator on meta.
251        // TODO: add some validation logic here.
252        Ok(vec![])
253    }
254}
255impl IcebergSplitEnumerator {
256    pub fn new_inner(properties: IcebergProperties, _context: SourceEnumeratorContextRef) -> Self {
257        Self { config: properties }
258    }
259}
260
261#[derive(Debug, Clone, PartialEq, Eq, Hash)]
262pub enum IcebergTimeTravelInfo {
263    Version(i64),
264    TimestampMs(i64),
265}
266
267#[derive(Debug, Clone)]
268pub struct IcebergListResult {
269    pub data_files: Vec<FileScanTask>,
270    pub equality_delete_files: Vec<FileScanTask>,
271    pub position_delete_files: Vec<FileScanTask>,
272    pub equality_delete_columns: Vec<String>,
273    pub format_version: FormatVersion,
274    pub schema: std::sync::Arc<iceberg::spec::Schema>,
275}
276
277impl IcebergSplitEnumerator {
278    pub fn get_snapshot_id(
279        table: &Table,
280        time_travel_info: Option<IcebergTimeTravelInfo>,
281    ) -> ConnectorResult<Option<i64>> {
282        let current_snapshot = table.metadata().current_snapshot();
283        let Some(current_snapshot) = current_snapshot else {
284            return Ok(None);
285        };
286
287        let snapshot_id = match time_travel_info {
288            Some(IcebergTimeTravelInfo::Version(version)) => {
289                let Some(snapshot) = table.metadata().snapshot_by_id(version) else {
290                    bail!("Cannot find the snapshot id in the iceberg table.");
291                };
292                snapshot.snapshot_id()
293            }
294            Some(IcebergTimeTravelInfo::TimestampMs(timestamp)) => {
295                let snapshot = table
296                    .metadata()
297                    .snapshots()
298                    .filter(|snapshot| snapshot.timestamp_ms() <= timestamp)
299                    .max_by_key(|snapshot| snapshot.timestamp_ms());
300                match snapshot {
301                    Some(snapshot) => snapshot.snapshot_id(),
302                    None => {
303                        // convert unix time to human-readable time
304                        let time = chrono::DateTime::from_timestamp_millis(timestamp);
305                        if let Some(time) = time {
306                            tracing::warn!("Cannot find a snapshot older than {}", time);
307                        } else {
308                            tracing::warn!("Cannot find a snapshot");
309                        }
310                        return Ok(None);
311                    }
312                }
313            }
314            None => current_snapshot.snapshot_id(),
315        };
316        Ok(Some(snapshot_id))
317    }
318
319    pub async fn list_scan_tasks(
320        &self,
321        time_travel_info: Option<IcebergTimeTravelInfo>,
322        predicate: IcebergPredicate,
323    ) -> ConnectorResult<Option<IcebergListResult>> {
324        let table = self.config.load_table().await?;
325        let snapshot_id = Self::get_snapshot_id(&table, time_travel_info)?;
326
327        let Some(snapshot_id) = snapshot_id else {
328            return Ok(None);
329        };
330        let res = self
331            .list_scan_tasks_inner(&table, snapshot_id, predicate)
332            .await?;
333        Ok(Some(res))
334    }
335
336    async fn list_scan_tasks_inner(
337        &self,
338        table: &Table,
339        snapshot_id: i64,
340        predicate: IcebergPredicate,
341    ) -> ConnectorResult<IcebergListResult> {
342        let format_version = table.metadata().format_version();
343        let table_schema = table.metadata().current_schema();
344        tracing::debug!("iceberg_table_schema: {:?}", table_schema);
345
346        let mut position_delete_files = vec![];
347        let mut position_delete_files_set = HashSet::new();
348        let mut data_files = vec![];
349        let mut equality_delete_files = vec![];
350        let mut equality_delete_files_set = HashSet::new();
351        let mut equality_delete_ids = None;
352        let mut scan_builder = table.scan().snapshot_id(snapshot_id).select_all();
353        if predicate != IcebergPredicate::AlwaysTrue {
354            scan_builder = scan_builder.with_filter(predicate.clone());
355        }
356        let scan = scan_builder.build()?;
357        let file_scan_stream = scan.plan_files().await?;
358
359        #[for_await]
360        for task in file_scan_stream {
361            let task: FileScanTask = task?;
362
363            // Collect delete files for separate scan types, but keep task.deletes intact
364            for delete_file in &task.deletes {
365                let delete_file = delete_file.as_ref().clone();
366                match delete_file.data_file_content {
367                    iceberg::spec::DataContentType::Data => {
368                        bail!("Data file should not in task deletes");
369                    }
370                    iceberg::spec::DataContentType::EqualityDeletes => {
371                        if equality_delete_files_set.insert(delete_file.data_file_path.clone()) {
372                            if equality_delete_ids.is_none() {
373                                equality_delete_ids = delete_file.equality_ids.clone();
374                            } else if equality_delete_ids != delete_file.equality_ids {
375                                bail!(
376                                    "The schema of iceberg equality delete file must be consistent"
377                                );
378                            }
379                            equality_delete_files.push(delete_file);
380                        }
381                    }
382                    iceberg::spec::DataContentType::PositionDeletes => {
383                        if position_delete_files_set.insert(delete_file.data_file_path.clone()) {
384                            position_delete_files.push(delete_file);
385                        }
386                    }
387                }
388            }
389
390            match task.data_file_content {
391                iceberg::spec::DataContentType::Data => {
392                    // Keep the original task with its deletes field intact
393                    data_files.push(task);
394                }
395                iceberg::spec::DataContentType::EqualityDeletes => {
396                    bail!("Equality delete files should not be in the data files");
397                }
398                iceberg::spec::DataContentType::PositionDeletes => {
399                    bail!("Position delete files should not be in the data files");
400                }
401            }
402        }
403        let schema = table_schema.clone();
404        let equality_delete_columns = equality_delete_ids
405            .unwrap_or_default()
406            .into_iter()
407            .map(|id| match schema.name_by_field_id(id) {
408                Some(name) => Ok::<std::string::String, ConnectorError>(name.to_owned()),
409                None => bail!("Delete field id {} not found in schema", id),
410            })
411            .collect::<ConnectorResult<Vec<_>>>()?;
412
413        Ok(IcebergListResult {
414            data_files,
415            equality_delete_files,
416            position_delete_files,
417            equality_delete_columns,
418            format_version,
419            schema,
420        })
421    }
422
423    /// Uniformly distribute scan tasks to compute nodes.
424    /// It's deterministic so that it can best utilize the data locality.
425    ///
426    /// # Arguments
427    /// * `file_scan_tasks`: The file scan tasks to be split.
428    /// * `split_num`: The number of splits to be created.
429    ///
430    /// 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.
431    /// Ensure that the total length of each group is as balanced as possible.
432    /// The time complexity is O(n log k), where n is the number of file scan tasks and k is the number of splits.
433    /// The space complexity is O(k), where k is the number of splits.
434    /// The algorithm is stable, so the order of the file scan tasks will be preserved.
435    pub fn split_n_vecs(
436        file_scan_tasks: Vec<FileScanTask>,
437        split_num: usize,
438    ) -> Vec<Vec<FileScanTask>> {
439        IcebergScanTaskPlanner::split_n_vecs(file_scan_tasks, split_num)
440    }
441}
442
443pub struct IcebergScanOpts {
444    pub chunk_size: usize,
445    pub need_seq_num: bool,
446    pub need_file_path_and_pos: bool,
447    pub handle_delete_files: bool,
448}
449
450/// Scan a data file. Delete files are handled by the iceberg-rust `reader.read` implementation.
451#[try_stream(ok = DataChunk, error = ConnectorError)]
452pub async fn scan_task_to_chunk_with_deletes(
453    table: Table,
454    mut data_file_scan_task: FileScanTask,
455    IcebergScanOpts {
456        chunk_size,
457        need_seq_num,
458        need_file_path_and_pos,
459        handle_delete_files,
460    }: IcebergScanOpts,
461    metrics: Option<Arc<IcebergScanMetrics>>,
462) {
463    let table_name = table.identifier().name().to_owned();
464
465    let num_delete_files = data_file_scan_task.deletes.len();
466    let expected_record_count = data_file_scan_task.record_count;
467    let file_start = std::time::Instant::now();
468
469    let mut read_bytes = scopeguard::guard(0u64, |read_bytes| {
470        if let Some(metrics) = metrics.clone() {
471            metrics
472                .iceberg_read_bytes
473                .with_guarded_label_values(&[&table_name])
474                .inc_by(read_bytes as _);
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));
500
501    let record_batch_stream: iceberg::scan::ArrowRecordBatchStream =
502        reader.read(Box::pin(file_scan_stream))?;
503    let mut record_batch_stream = record_batch_stream.enumerate();
504
505    let mut total_rows_read: u64 = 0;
506
507    // Process each record batch. Delete application is handled by the SDK.
508    while let Some((batch_index, record_batch)) = record_batch_stream.next().await {
509        let record_batch = record_batch?;
510        let batch_start_pos = (batch_index * chunk_size) as i64;
511
512        let mut chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
513        let row_count = chunk.capacity();
514        total_rows_read += row_count as u64;
515
516        // Add metadata columns if requested
517        if need_seq_num {
518            let (mut columns, visibility) = chunk.into_parts();
519            columns.push(Arc::new(ArrayImpl::Int64(I64Array::from_iter(
520                std::iter::repeat_n(data_sequence_number, row_count),
521            ))));
522            chunk = DataChunk::from_parts(columns.into(), visibility);
523        }
524
525        if need_file_path_and_pos {
526            let (mut columns, visibility) = chunk.into_parts();
527            columns.push(Arc::new(ArrayImpl::Utf8(Utf8Array::from_iter(
528                std::iter::repeat_n(data_file_path.as_str(), row_count),
529            ))));
530
531            // Generate position values for each row in the batch
532            let positions: Vec<i64> =
533                (batch_start_pos..(batch_start_pos + row_count as i64)).collect();
534            columns.push(Arc::new(ArrayImpl::Int64(I64Array::from_iter(positions))));
535
536            chunk = DataChunk::from_parts(columns.into(), visibility);
537        }
538
539        *read_bytes += chunk.estimated_heap_size() as u64;
540        yield chunk;
541    }
542
543    // Record per-file metrics after reading all batches.
544    if let Some(ref metrics) = metrics {
545        let label_values = [table_name.as_str()];
546
547        // File read duration.
548        metrics
549            .iceberg_source_file_read_duration_seconds
550            .with_guarded_label_values(&label_values)
551            .observe(file_start.elapsed().as_secs_f64());
552
553        // Rows read.
554        if total_rows_read > 0 {
555            metrics
556                .iceberg_source_rows_read_total
557                .with_guarded_label_values(&label_values)
558                .inc_by(total_rows_read);
559        }
560
561        // File read count.
562        metrics
563            .iceberg_source_files_read_total
564            .with_guarded_label_values(&[table_name.as_str(), "data"])
565            .inc();
566
567        // APPROXIMATE: Estimate delete rows applied. The delta between expected_record_count
568        // and actual rows read may also include predicate pushdown / row-group pruning effects,
569        // so this metric can overcount. It is still useful as an approximate signal for
570        // detecting whether delete files cause significant row filtering.
571        if handle_delete_files
572            && num_delete_files > 0
573            && let Some(expected) = expected_record_count
574        {
575            let deleted = expected.saturating_sub(total_rows_read);
576            if deleted > 0 {
577                metrics
578                    .iceberg_source_delete_rows_applied_total
579                    .with_guarded_label_values(&[table_name.as_str(), "sdk_applied_approx"])
580                    .inc_by(deleted);
581            }
582        }
583    }
584}
585
586#[derive(Debug)]
587pub struct IcebergFileReader {}
588
589#[async_trait]
590impl SplitReader for IcebergFileReader {
591    type Properties = IcebergProperties;
592    type Split = IcebergSplit;
593
594    async fn new(
595        _props: IcebergProperties,
596        _splits: Vec<IcebergSplit>,
597        _parser_config: ParserConfig,
598        _source_ctx: SourceContextRef,
599        _columns: Option<Vec<Column>>,
600    ) -> ConnectorResult<Self> {
601        unimplemented!()
602    }
603
604    fn into_stream(self) -> BoxSourceChunkStream {
605        unimplemented!()
606    }
607}
608
609#[cfg(test)]
610mod tests {
611    use std::sync::Arc;
612
613    use iceberg::scan::FileScanTask;
614    use iceberg::spec::{DataContentType, Schema};
615
616    use super::*;
617
618    fn create_file_scan_task(length: u64, id: u64) -> FileScanTask {
619        FileScanTask {
620            length,
621            start: 0,
622            record_count: Some(0),
623            data_file_path: format!("test_{}.parquet", id),
624            referenced_data_file: None,
625            data_file_content: DataContentType::Data,
626            data_file_format: iceberg::spec::DataFileFormat::Parquet,
627            schema: Arc::new(Schema::builder().build().unwrap()),
628            project_field_ids: vec![],
629            predicate: None,
630            deletes: vec![],
631            sequence_number: 0,
632            equality_ids: None,
633            file_size_in_bytes: 0,
634            partition: None,
635            partition_spec: None,
636            name_mapping: None,
637            case_sensitive: true,
638        }
639    }
640
641    #[test]
642    fn test_split_n_vecs_basic() {
643        let file_scan_tasks = (1..=12)
644            .map(|i| create_file_scan_task(i + 100, i))
645            .collect::<Vec<_>>(); // Ensure the correct function is called
646
647        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
648
649        assert_eq!(groups.len(), 3);
650
651        let group_lengths: Vec<u64> = groups
652            .iter()
653            .map(|group| group.iter().map(|task| task.length).sum())
654            .collect();
655
656        let max_length = *group_lengths.iter().max().unwrap();
657        let min_length = *group_lengths.iter().min().unwrap();
658        assert!(max_length - min_length <= 10, "Groups should be balanced");
659
660        let total_tasks: usize = groups.iter().map(|group| group.len()).sum();
661        assert_eq!(total_tasks, 12);
662    }
663
664    #[test]
665    fn test_split_n_vecs_empty() {
666        let file_scan_tasks = Vec::new();
667        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
668        assert_eq!(groups.len(), 3);
669        assert!(groups.iter().all(|group| group.is_empty()));
670    }
671
672    #[test]
673    fn test_split_n_vecs_single_task() {
674        let file_scan_tasks = vec![create_file_scan_task(100, 1)];
675        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 3);
676        assert_eq!(groups.len(), 3);
677        assert_eq!(groups.iter().filter(|group| !group.is_empty()).count(), 1);
678    }
679
680    #[test]
681    fn test_split_n_vecs_uneven_distribution() {
682        let file_scan_tasks = vec![
683            create_file_scan_task(1000, 1),
684            create_file_scan_task(100, 2),
685            create_file_scan_task(100, 3),
686            create_file_scan_task(100, 4),
687            create_file_scan_task(100, 5),
688        ];
689
690        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks, 2);
691        assert_eq!(groups.len(), 2);
692
693        let group_with_large_task = groups
694            .iter()
695            .find(|group| group.iter().any(|task| task.length == 1000))
696            .unwrap();
697        assert_eq!(group_with_large_task.len(), 1);
698    }
699
700    #[test]
701    fn test_split_n_vecs_same_files_distribution() {
702        let file_scan_tasks = vec![
703            create_file_scan_task(100, 1),
704            create_file_scan_task(100, 2),
705            create_file_scan_task(100, 3),
706            create_file_scan_task(100, 4),
707            create_file_scan_task(100, 5),
708            create_file_scan_task(100, 6),
709            create_file_scan_task(100, 7),
710            create_file_scan_task(100, 8),
711        ];
712
713        let groups = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks.clone(), 4)
714            .iter()
715            .map(|g| {
716                g.iter()
717                    .map(|task| task.data_file_path.clone())
718                    .collect::<Vec<_>>()
719            })
720            .collect::<Vec<_>>();
721
722        for _ in 0..10000 {
723            let groups_2 = IcebergSplitEnumerator::split_n_vecs(file_scan_tasks.clone(), 4)
724                .iter()
725                .map(|g| {
726                    g.iter()
727                        .map(|task| task.data_file_path.clone())
728                        .collect::<Vec<_>>()
729                })
730                .collect::<Vec<_>>();
731
732            assert_eq!(groups, groups_2);
733        }
734    }
735}
risingwave_connector/source/iceberg/mod.rs

risingwave_connector/source/iceberg/
mod.rs
