feat(iceberg): support eq delete merge on read for iceberg

Cargo.toml

@@ -142,9 +142,9 @@ arrow-array-iceberg = { package = "arrow-array", version = "52" }
 arrow-schema-iceberg = { package = "arrow-schema", version = "52" }
 arrow-buffer-iceberg = { package = "arrow-buffer", version = "52" }
 arrow-cast-iceberg = { package = "arrow-cast", version = "52" }
-iceberg = "0.3.0"
-iceberg-catalog-rest = "0.3.0"
-iceberg-catalog-glue = "0.3.0"
+iceberg = { git = "https://github.com/risingwavelabs/iceberg-rust.git", rev = "2648421c18993edec6853ad5ce978733178b26f1" }
+iceberg-catalog-rest = { git = "https://github.com/risingwavelabs/iceberg-rust.git", rev = "2648421c18993edec6853ad5ce978733178b26f1" }
+iceberg-catalog-glue = { git = "https://github.com/risingwavelabs/iceberg-rust.git", rev = "2648421c18993edec6853ad5ce978733178b26f1" }
 opendal = "0.47"
 arrow-array = "50"
 arrow-arith = "50"

src/batch/src/executor/iceberg_scan.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use core::ops::BitOr;
+use std::collections::HashMap;
 use std::mem;
 
 use futures_async_stream::try_stream;
@@ -20,6 +22,7 @@ use iceberg::scan::FileScanTask;
 use iceberg::spec::TableMetadata;
 use itertools::Itertools;
 use risingwave_common::array::arrow::IcebergArrowConvert;
+use risingwave_common::bitmap::Bitmap;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::types::DataType;
 use risingwave_connector::sink::iceberg::IcebergConfig;
@@ -39,6 +42,7 @@ pub struct IcebergScanExecutor {
     snapshot_id: Option<i64>,
     table_meta: TableMetadata,
     file_scan_tasks: Vec<FileScanTask>,
+    eq_delete_file_scan_tasks: Vec<FileScanTask>,
     batch_size: usize,
     schema: Schema,
     identity: String,
@@ -64,6 +68,7 @@ impl IcebergScanExecutor {
         snapshot_id: Option<i64>,
         table_meta: TableMetadata,
         file_scan_tasks: Vec<FileScanTask>,
+        eq_delete_file_scan_tasks: Vec<FileScanTask>,
         batch_size: usize,
         schema: Schema,
         identity: String,
@@ -73,6 +78,7 @@ impl IcebergScanExecutor {
             snapshot_id,
             table_meta,
             file_scan_tasks,
+            eq_delete_file_scan_tasks,
             batch_size,
             schema,
             identity,
@@ -86,33 +92,122 @@ impl IcebergScanExecutor {
             .load_table_v2_with_metadata(self.table_meta)
             .await?;
         let data_types = self.schema.data_types();
+        let chunk_schema_names = self.schema.names();
 
-        let file_scan_tasks = mem::take(&mut self.file_scan_tasks);
+        let mut eq_delete_file_scan_tasks_map: HashMap<
+            String,
+            HashMap<Option<risingwave_common::types::ScalarImpl>, i64>,
+        > = HashMap::default();
+        let eq_delete_file_scan_tasks = mem::take(&mut self.eq_delete_file_scan_tasks);
+
+        for eq_delete_file_scan_task in eq_delete_file_scan_tasks {
+            let mut sequence_number = eq_delete_file_scan_task.sequence_number();
+            let reader = table
+                .reader_builder()
+                .with_batch_size(self.batch_size)
+                .build();
+            let delete_file_scan_stream = tokio_stream::once(Ok(eq_delete_file_scan_task));
+
+            let mut delete_record_batch_stream = reader
+                .read(Box::pin(delete_file_scan_stream))
+                .map_err(BatchError::Iceberg)?;
 
-        let file_scan_stream = {
-            #[try_stream]
-            async move {
-                for file_scan_task in file_scan_tasks {
-                    yield file_scan_task;
+            while let Some(record_batch) = delete_record_batch_stream.next().await {
+                let record_batch = record_batch.map_err(BatchError::Iceberg)?;
+                let delete_column_names = record_batch
+                    .schema()
+                    .fields()
+                    .iter()
+                    .map(|field| field.name())
+                    .cloned()
+                    .collect_vec();
+                let chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
+                for (array, columns_name) in chunk.columns().iter().zip_eq(delete_column_names) {
+                    let each_column_seq_num_map = eq_delete_file_scan_tasks_map
+                        .entry(columns_name)
+                        .or_default();
+                    for datum in array.get_all_values() {
+                        let entry = each_column_seq_num_map
+                            .entry(datum)
+                            .or_insert(sequence_number);
+                        *entry = *entry.max(&mut sequence_number);
+                    }
                 }
             }
-        };
-
-        let reader = table
-            .reader_builder()
-            .with_batch_size(self.batch_size)
-            .build();
-
-        let record_batch_stream = reader
-            .read(Box::pin(file_scan_stream))
-            .map_err(BatchError::Iceberg)?;
-
-        #[for_await]
-        for record_batch in record_batch_stream {
-            let record_batch = record_batch.map_err(BatchError::Iceberg)?;
-            let chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
-            debug_assert_eq!(chunk.data_types(), data_types);
-            yield chunk;
+        }
+
+        let file_scan_tasks = mem::take(&mut self.file_scan_tasks);
+
+        for file_scan_task in file_scan_tasks {
+            let sequence_number = file_scan_task.sequence_number();
+            let reader = table
+                .reader_builder()
+                .with_batch_size(self.batch_size)
+                .build();
+            let file_scan_stream = tokio_stream::once(Ok(file_scan_task));
+
+            let mut record_batch_stream = reader
+                .read(Box::pin(file_scan_stream))
+                .map_err(BatchError::Iceberg)?;
+
+            while let Some(record_batch) = record_batch_stream.next().await {
+                let record_batch = record_batch.map_err(BatchError::Iceberg)?;
+                let column_names = record_batch
+                    .schema()
+                    .fields()
+                    .iter()
+                    .map(|field| field.name())
+                    .cloned()
+                    .collect_vec();
+                let chunk = IcebergArrowConvert.chunk_from_record_batch(&record_batch)?;
+                let visibilitys: Vec<_> = chunk
+                    .columns()
+                    .iter()
+                    .zip_eq(column_names.clone())
+                    .filter_map(|(array, column_map)| {
+                        if let Some(each_column_seq_num_map) =
+                            eq_delete_file_scan_tasks_map.get(&column_map)
+                        {
+                            let visibility =
+                                Bitmap::from_iter(array.get_all_values().iter().map(|datum| {
+                                    if let Some(s) = each_column_seq_num_map.get(datum)
+                                        && s > &sequence_number
+                                    {
+                                        false
+                                    } else {
+                                        true
+                                    }
+                                }));
+                            Some(visibility)
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+                let (data, va) = chunk.into_parts_v2();
+                let visibility = if visibilitys.is_empty() {
+                    va
+                } else {
+                    visibilitys
+                        .iter()
+                        .skip(1)
+                        .fold(visibilitys[0].clone(), |acc, bitmap| acc.bitor(bitmap))
+                };
+                let data = data
+                    .iter()
+                    .zip_eq(column_names)
+                    .filter_map(|(array, columns)| {
+                        if chunk_schema_names.contains(&columns) {
+                            Some(array.clone())
+                        } else {
+                            None
+                        }
+                    })
+                    .collect_vec();
+                let chunk = DataChunk::new(data, visibility);
+                debug_assert_eq!(chunk.data_types(), data_types);
+                yield chunk;
+            }
         }
     }
 }
@@ -171,6 +266,11 @@ impl BoxedExecutorBuilder for IcebergScanExecutorBuilder {
                 Some(split.snapshot_id),
                 split.table_meta.deserialize(),
                 split.files.into_iter().map(|x| x.deserialize()).collect(),
+                split
+                    .eq_delete_files
+                    .into_iter()
+                    .map(|x| x.deserialize())
+                    .collect(),
                 source.context.get_config().developer.chunk_size,
                 schema,
                 source.plan_node().get_identity().clone(),

src/common/src/array/mod.rs

@@ -603,6 +603,10 @@ impl ArrayImpl {
         })
     }
 
+    pub fn get_all_values(&self) -> Vec<Datum> {
+        (0..self.len()).map(|i| self.datum_at(i)).collect()
+    }
+
     /// # Safety
     ///
     /// This function is unsafe because it does not check the validity of `idx`. It is caller's

src/connector/src/source/iceberg/mod.rs

@@ -14,13 +14,16 @@
 
 pub mod parquet_file_reader;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
 
 use anyhow::anyhow;
 use async_trait::async_trait;
 use futures_async_stream::for_await;
+use iceberg::arrow::schema_to_arrow_schema;
 use iceberg::scan::FileScanTask;
 use iceberg::spec::TableMetadata;
+use iceberg::table::Table;
 use itertools::Itertools;
 pub use parquet_file_reader::*;
 use risingwave_common::bail;
@@ -144,6 +147,7 @@ pub struct IcebergSplit {
     pub snapshot_id: i64,
     pub table_meta: TableMetadataJsonStr,
     pub files: Vec<IcebergFileScanTaskJsonStr>,
+    pub eq_delete_files: Vec<IcebergFileScanTaskJsonStr>,
 }
 
 impl SplitMetaData for IcebergSplit {
@@ -237,12 +241,23 @@ impl IcebergSplitEnumerator {
                 None => bail!("Cannot find the current snapshot id in the iceberg table."),
             },
         };
+        let (eq_delete_files, eq_delete_file_schema) =
+            IcebergSplitEnumerator::load_eq_delete_file(&table, snapshot_id).await?;
+        let arrow_schema = schema_to_arrow_schema(&eq_delete_file_schema)?;
+        let mut require_names: HashSet<String> = schema.names().clone().into_iter().collect();
+        require_names.extend(
+            arrow_schema
+                .all_fields()
+                .into_iter()
+                .map(|filed| filed.name().clone()),
+        );
+
         let mut files = vec![];
 
         let scan = table
             .scan()
             .snapshot_id(snapshot_id)
-            .select(schema.names())
+            .select(require_names)
             .build()
             .map_err(|e| anyhow!(e))?;
 
@@ -269,6 +284,7 @@ impl IcebergSplitEnumerator {
                 snapshot_id,
                 table_meta: table_meta.clone(),
                 files: files[start..end].to_vec(),
+                eq_delete_files: eq_delete_files.clone(),
             };
             splits.push(split);
         }
@@ -282,6 +298,29 @@ impl IcebergSplitEnumerator {
             .filter(|split| !split.files.is_empty())
             .collect_vec())
     }
+
+    async fn load_eq_delete_file(
+        table: &Table,
+        snapshot_id: i64,
+    ) -> ConnectorResult<(Vec<IcebergFileScanTaskJsonStr>, Arc<iceberg::spec::Schema>)> {
+        let mut files = vec![];
+
+        let scan = table
+            .scan()
+            .snapshot_id(snapshot_id)
+            .build()
+            .map_err(|e| anyhow!(e))?;
+        let schema = scan.snapshot().schema(table.metadata())?;
+
+        let file_scan_stream = scan.plan_eq_delete_files().await.map_err(|e| anyhow!(e))?;
+
+        #[for_await]
+        for task in file_scan_stream {
+            let task = task.map_err(|e| anyhow!(e))?;
+            files.push(IcebergFileScanTaskJsonStr::serialize(&task));
+        }
+        Ok((files, schema))
+    }
 }
 
 #[derive(Debug)]